nemo_automodel.components.models.minimax_m3_vl.layers

Bases: Module

MiniMax M3 decoder block: attention + (dense MLP or MoE) with Gemma norms.

moe_layer_freq[layer_idx] == 0 -> dense MiniMaxM3MLP (with dense_intermediate_size); otherwise a routed MoE plus a separate SwiGLU-OAI shared expert (kept M3-local rather than using MoE’s built-in shared expert, whose generic MLP does not implement SwiGLU-OAI).

nemo_automodel.components.models.minimax_m3_vl.layers.Block.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

nemo_automodel.components.models.minimax_m3_vl.layers.Block.init_weights(
    buffer_device: torch.device
)

Bases: Module

MiniMax M3 GQA attention with per-head Gemma Q/K norm and partial RoPE.

When is_sparse_attention_layer is set, an additional lightning indexer (index_q/k_proj + per-head Gemma norm) selects, per query, the top-k key blocks to attend to (block-level DeepSeek-style sparse attention). M3 sets disable_index_value=True so the index branch is selection-only.

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Attention.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Attention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
)

Bases: Module

Lightning indexer (selection-only) for MiniMax M3 sparse-attention layers.

Projects hidden states to num_index_heads index queries and a single shared index key (disable_index_value=True for M3, so there is no index value/output projection). Per-head Gemma RMSNorm + partial RoPE mirror the main attention. The produced idx_q/idx_k feed :func:build_block_sparse_attn_bias to select which key blocks each query attends to.

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Indexer.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    num_q_heads: int,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Indexer.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
)

Bases: Module

Dense / shared-expert MLP with SwiGLU-OAI activation (separate gate/up/down).

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3MLP.forward(
    x: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3MLP.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Bases: Module

RMSNorm with optional Gemma-style zero-centered gamma (x_normed * (1 + w)).

When gemma=True the learnable weight is centered at 0 and the effective scale is 1 + weight (matching HF GemmaRMSNorm and the sglang M3 reference). Used both for hidden-size norms and, with dim=head_dim, for per-head Q/K normalization (the input is normalized over its last dim, so a [..., num_heads, head_dim] tensor is normalized independently per head).

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3RMSNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3RMSNorm.reset_parameters() -> None

Convert an incoming padding mask to a boolean key keep-mask broadcastable to ref.

Accepts a 2-D [B, T] keep-mask (1/True = attend) or an already-boolean 4-D mask; returns a boolean mask (True where the key is attendable) to be AND-ed with the block-sparse keep-mask. Boolean (not additive) so the eager SDPA path is bf16-safe — see :func:build_block_sparse_attn_mask.

Block selection for one query chunk (see :func:select_sparse_blocks).

Build the boolean [B, num_q_heads, T, T] block-sparse causal keep-mask.

Eager (i.e non-CP) path: selects blocks via :func:select_sparse_blocks over the square sequence, expands the block selection to a per-key mask, intersects with token-level causal, and returns a boolean keep-mask (True where attended) repeat-interleaved across GQA groups.

NOTE: returns a boolean mask, NOT an additive 0/-inf bias. An additive -inf bias is numerically unsafe under SDPA in bf16 — it leaks past the mask at early (few-key) positions, while finfo.min produces NaNs. SDPA masks correctly with a boolean attn_mask (matching the CP path’s FlexAttention BlockMask).

Parameters:

Select, per query, which key blocks to attend to (DSA block top-k).

Mirrors the sglang minimax_sparse selection (block_size_q=1 -> per-query-position): the index score for (query i, key j) is (idx_q[i] . idx_k[j]) * idx_dim**-0.5 with causal masking; keys are grouped into blocks of block_size and reduced per block (max or lse). For each query, the current block (local_blocks) and the first init_blocks are always kept and the remaining budget is filled with the highest-scoring causal blocks, up to min(topk_blocks, valid_blocks).

Queries and keys are decoupled so the same selection serves the eager square case (Tq == Tk, q_positions defaulting to arange) and the context-parallel case (local queries Tq against the gathered global key sequence Tk, with q_positions giving each local query’s global position). Key block b spans key indices [b*block_size, ...) in the (global) key order. The query dim is chunked so the fp32 score tensor stays within _SELECT_SCORE_BUDGET_BYTES (the long-context memory bottleneck); per-query independence makes chunking exact (concatenated along Tq).

Parameters:

Returns: torch.Tensor

[B, H_idx, Tq, num_blocks] bool block-selection mask (causal +

GPT-OSS / MiniMax-M3 SwiGLU-OAI: gate * sigmoid(alpha * gate) * (up + 1).

Gate is clamped max=limit and up is clamped +/-limit (when limit > 0), computed in fp32 and cast back. Equivalent to sglang’s swiglu_no_interleaved_with_alpha_and_limit.