mstar.model.components.aut_encoder

AuT: the Qwen3-family audio encoder (Qwen3-Omni, Qwen3-ASR), natively.

A Whisper-style log-mel input (128 bins, 100 frames/s) is cut into 1 s chunks of 2 * n_window frames, each chunk runs through three stride-2 Conv2d layers (100 frames -> 13 tokens, so ~13 tokens/s) and a linear conv_out, gets the same 13 sinusoidal positions, and the valid tokens of every chunk are packed back into one [total_tokens, d_model] sequence. The transformer stack then attends within windows of n_window_infer frames (8 s -> 104 tokens): block-diagonal, bidirectional, no cache. A final LayerNorm and a two-layer GELU projector map the tokens into the LLM’s width.

Windows are the attention’s only structure, so the stack is packed varlen attention: every window is a segment, and any number of requests can be packed into one forward. On the engine it runs through the cacheless RaggedAttentionSpec resource, whose plan the runner builds outside the graph — which is what lets the block loop replay as a CUDA graph. Without a bound resource (CPU tests, standalone use) it falls back to SDPA per window.

Parameter paths mirror HF’s audio_tower.* (Qwen3OmniMoeAudioEncoder, which Qwen3-ASR reuses unchanged), with q/k/v_proj fused into qkv_proj by the loader’s Whisper stacked rules.

Classes

AuTAttention(d_model, num_heads, attn_key)	Bidirectional attention within the packed windows of one forward.
AuTEncoder(config[, attn_key])	Packed audio encoder; parameter paths mirror HF's audio_tower.*.
AuTEncoderConfig([d_model, encoder_layers, ...])
AuTEncoderLayer(config, attn_key)
AuTLayout(tokens_per_request, cu_seqlens)	Where one packed forward's tokens come from and how they attend.

class mstar.model.components.aut_encoder.AuTAttention(d_model, num_heads, attn_key)

Bases: Module

Bidirectional attention within the packed windows of one forward.

attn_key names the node’s RaggedAttentionSpec resource; bound at load, it attends through the runner’s plan. Unbound, SDPA runs per window from cu_seqlens.

Parameters:

• d_model (int)

• num_heads (int)

• attn_key (str | None)

bind_resources(resources)

Parameters:

resources (dict)

Return type:

None

forward(hidden_states, cu_seqlens)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.aut_encoder.AuTEncoder(config, attn_key=None)

Bases: Module

Packed audio encoder; parameter paths mirror HF’s audio_tower.*.

Parameters:

• config (AuTEncoderConfig)

• attn_key (str | None)

encode(hidden_states, cu_seqlens)

The transformer stack and projector over a packed sequence: (total_tokens, d_model) -> (total_tokens, output_dim). This is the region a submodule captures.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

forward(features, feature_lens)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• features (Tensor)

• feature_lens (list[int])

Return type:

tuple[Tensor, AuTLayout]

frontend(features, feature_lens)

(B, num_mel_bins, T_max) zero-padded mel + valid lengths -> packed (total_tokens, d_model) chunk embeddings with positions.

Every request is cut into chunk_frames pieces; the last piece is zero-padded to a full chunk (as HF’s pad_sequence does whenever a batch holds one full chunk), convolved, and trimmed back to its valid tokens.

Parameters:

• features (Tensor)

• feature_lens (list[int])

Return type:

Tensor

layout(feature_lens)

Token counts and attention windows for a batch of mel lengths. Host-side arithmetic only, so a submodule can declare its step from it.

Parameters:

feature_lens (list[int])

Return type:

AuTLayout

load_weights(weights)

Stream audio_tower.*-relative (name, tensor) pairs in; raise if any parameter is left unloaded.

Return type:

set[str]

reset_buffers()

Recompute the (non-persistent) sinusoid table. Needed after a meta-device construction + to_empty, which leaves buffers uninitialized; load_weights calls it.

Return type:

None

class mstar.model.components.aut_encoder.AuTEncoderConfig(d_model: 'int' = 1024, encoder_layers: 'int' = 24, encoder_attention_heads: 'int' = 16, encoder_ffn_dim: 'int' = 4096, num_mel_bins: 'int' = 128, n_window: 'int' = 50, n_window_infer: 'int' = 800, conv_chunksize: 'int' = 500, downsample_hidden_size: 'int' = 480, output_dim: 'int' = 2048, max_source_positions: 'int' = 1500, activation_function: 'str' = 'gelu')

Bases: object

Parameters:

• d_model (int)

• encoder_layers (int)

• encoder_attention_heads (int)

• encoder_ffn_dim (int)

• num_mel_bins (int)

• n_window (int)

• n_window_infer (int)

• conv_chunksize (int)

• downsample_hidden_size (int)

• output_dim (int)

• max_source_positions (int)

• activation_function (str)

activation_function: str = 'gelu'

property chunk_frames: int

conv_chunksize: int = 500

property conv_out_features: int

d_model: int = 1024

downsample_hidden_size: int = 480

encoder_attention_heads: int = 16

encoder_ffn_dim: int = 4096

encoder_layers: int = 24

classmethod from_hf(hf)

Parameters:

hf (dict)

Return type:

AuTEncoderConfig

property head_dim: int

max_source_positions: int = 1500

n_window: int = 50

n_window_infer: int = 800

num_mel_bins: int = 128

output_dim: int = 2048

static tokens_for_frames(num_frames)

Encoder tokens for num_frames mel frames: 13 per full 100-frame chunk, plus three stride-2 convolutions’ worth of the remainder.

Parameters:

num_frames (int)

Return type:

int

property tokens_per_chunk: int

window_lengths(num_tokens)

Attention segments for one request’s tokens: full windows, then the remainder.

Parameters:

num_tokens (int)

Return type:

list[int]

property window_tokens: int

class mstar.model.components.aut_encoder.AuTEncoderLayer(config, attn_key)

Bases: Module

Parameters:

• config (AuTEncoderConfig)

• attn_key (str | None)

forward(hidden_states, cu_seqlens)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.aut_encoder.AuTLayout(tokens_per_request, cu_seqlens)

Bases: object

Where one packed forward’s tokens come from and how they attend.

Parameters:

• tokens_per_request (list[int])

• cu_seqlens (list[int])

cu_seqlens: list[int]

tokens_per_request: list[int]

property total_tokens: int

property window_lengths: list[int]