> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.minimax_m3_vl.layers

MiniMax M3 VL text-backbone layers.

Stage 1 covers the dense + MoE text path (no sparse-attention index branch and
no MTP).  Mirrors the canonical sglang reference
`sglang.srt.models.minimax_m3` (`MiniMaxM3Attention` / `MiniMaxM3MLP` /
`MiniMaxM3MoE` / `MiniMaxM3DecoderLayer`):

* per-head **Gemma** RMSNorm on Q/K (`qk_norm_type='per_head'`,
  `use_gemma_norm=True`),
* partial RoPE (`rotary_dim=64` of `head_dim=128`) reusing the gpt\_oss
  rotary utilities (as the existing `minimax_m2` backbone does),
* SwiGLU-OAI activation `gate * sigmoid(alpha * gate) * (up + 1)` with gate
  clamped `max=limit` and up clamped `+/-limit` for dense and shared experts,
* per-layer dense-vs-MoE selection from `moe_layer_freq`.

## Module Contents

### Classes

| Name                                                                                              | Description                                                                   |
| ------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`Block`](#nemo_automodel-components-models-minimax_m3_vl-layers-Block)                           | MiniMax M3 decoder block: attention + (dense MLP or MoE) with Gemma norms.    |
| [`MiniMaxM3Attention`](#nemo_automodel-components-models-minimax_m3_vl-layers-MiniMaxM3Attention) | MiniMax M3 GQA attention with per-head Gemma Q/K norm and partial RoPE.       |
| [`MiniMaxM3Indexer`](#nemo_automodel-components-models-minimax_m3_vl-layers-MiniMaxM3Indexer)     | Lightning indexer (selection-only) for MiniMax M3 sparse-attention layers.    |
| [`MiniMaxM3MLP`](#nemo_automodel-components-models-minimax_m3_vl-layers-MiniMaxM3MLP)             | Dense / shared-expert MLP with SwiGLU-OAI activation (separate gate/up/down). |
| [`MiniMaxM3RMSNorm`](#nemo_automodel-components-models-minimax_m3_vl-layers-MiniMaxM3RMSNorm)     | RMSNorm with optional Gemma-style zero-centered gamma (`x_normed * (1 + w)`). |

### Functions

| Name                                                                                                                  | Description                                                                         |
| --------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`_padding_mask_to_keep_mask`](#nemo_automodel-components-models-minimax_m3_vl-layers-_padding_mask_to_keep_mask)     | Convert an incoming padding mask to a boolean key keep-mask broadcastable to `ref`. |
| [`_select_sparse_blocks_chunk`](#nemo_automodel-components-models-minimax_m3_vl-layers-_select_sparse_blocks_chunk)   | Block selection for one query chunk (see :func:`select_sparse_blocks`).             |
| [`build_block_sparse_attn_mask`](#nemo_automodel-components-models-minimax_m3_vl-layers-build_block_sparse_attn_mask) | Build the boolean `[B, num_q_heads, T, T]` block-sparse causal keep-mask.           |
| [`select_sparse_blocks`](#nemo_automodel-components-models-minimax_m3_vl-layers-select_sparse_blocks)                 | Select, per query, which key blocks to attend to (DSA block top-k).                 |
| [`swiglu_oai`](#nemo_automodel-components-models-minimax_m3_vl-layers-swiglu_oai)                                     | GPT-OSS / MiniMax-M3 SwiGLU-OAI: `gate * sigmoid(alpha * gate) * (up + 1)`.         |

### Data

[`_SELECT_SCORE_BUDGET_BYTES`](#nemo_automodel-components-models-minimax_m3_vl-layers-_SELECT_SCORE_BUDGET_BYTES)

### API

```python
class nemo_automodel.components.models.minimax_m3_vl.layers.Block(
    layer_idx: int,
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**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).

```python
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
```

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

```python
class nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Attention(
    config: typing.Any,
    backend: nemo_automodel.components.models.common.BackendConfig,
    is_sparse_attention_layer: bool = False
)
```

**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.

```python
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
```

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

```python
class nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Indexer(
    config: typing.Any,
    sparse_cfg: dict,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**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.

```python
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
```

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

```python
class nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3MLP(
    config: typing.Any,
    intermediate_size: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

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

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

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

```python
class nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3RMSNorm(
    dim: int,
    eps: float = 1e-06,
    gemma: bool = True
)
```

**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).

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

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

```python
nemo_automodel.components.models.minimax_m3_vl.layers._padding_mask_to_keep_mask(
    attention_mask: torch.Tensor,
    ref: torch.Tensor
) -> torch.Tensor
```

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`.

```python
nemo_automodel.components.models.minimax_m3_vl.layers._select_sparse_blocks_chunk(
    idx_q: torch.Tensor,
    idx_k: torch.Tensor,
    q_positions: torch.Tensor,
    block_size: int,
    topk_blocks: int,
    init_blocks: int,
    local_blocks: int,
    score_type: str
) -> torch.Tensor
```

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

```python
nemo_automodel.components.models.minimax_m3_vl.layers.build_block_sparse_attn_mask(
    idx_q: torch.Tensor,
    idx_k: torch.Tensor,
    block_size: int,
    topk_blocks: int,
    init_blocks: int,
    local_blocks: int,
    num_q_heads: int,
    score_type: str = 'max'
) -> torch.Tensor
```

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:**

`[B, T, H_idx, D]` index queries (post norm + RoPE).

`[B, T, 1, D]` shared index key (post norm + RoPE).

number of main attention heads; the per-idx-head mask is
expanded `num_q_heads // H_idx` times (GQA, repeat-interleave).

```python
nemo_automodel.components.models.minimax_m3_vl.layers.select_sparse_blocks(
    idx_q: torch.Tensor,
    idx_k: torch.Tensor,
    block_size: int,
    topk_blocks: int,
    init_blocks: int,
    local_blocks: int,
    score_type: str = 'max',
    q_positions: torch.Tensor | None = None
) -> torch.Tensor
```

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:**

`[B, Tq, H_idx, D]` index queries (post norm + RoPE).

`[B, Tk, 1, D]` shared index key (post norm + RoPE).

`[Tq]` long global position of each query; defaults to
`arange(Tq)` (the eager square case).

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.models.minimax_m3_vl.layers.swiglu_oai(
    gate: torch.Tensor,
    up: torch.Tensor,
    alpha: float,
    limit: float
) -> torch.Tensor
```

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`.

```python
nemo_automodel.components.models.minimax_m3_vl.layers._SELECT_SCORE_BUDGET_BYTES = 2 * 1024 ** 3
```