> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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).     |
| [`MiniMaxM3MSAAttention`](#nemo_automodel-components-models-minimax_m3_vl-layers-MiniMaxM3MSAAttention) | A sparse attention layer running MSA's SM100 kernels on compact packed documents. |
| [`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                                                                         |
| --------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`_compiled_rms_norm`](#nemo_automodel-components-models-minimax_m3_vl-layers-_compiled_rms_norm)                     | Share one lazily compiled stateless core while keeping eager instances eager.       |
| [`_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`. |
| [`_rms_norm_core`](#nemo_automodel-components-models-minimax_m3_vl-layers-_rms_norm_core)                             | Evaluate M3's FP32 normalization and scale before the final activation cast.        |
| [`_select_sparse_blocks_chunk`](#nemo_automodel-components-models-minimax_m3_vl-layers-_select_sparse_blocks_chunk)   | Block selection for one query chunk (see `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).

**`input_layernorm`**

---

**`is_moe_layer`**

---

**`mlp`** `= MoE(moe_config, backend)`

---

**`post_attention_layernorm`**

---

**`self_attn`** `= MiniMaxM3MSAAttention(config, backend)`

---

**`shared_experts`** `= MiniMaxM3MLP(config, shared_inter, backend)`

---

```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
nemo_automodel.components.models.minimax_m3_vl.layers.Block.nemo_checkpoint_context_fn(
    context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
) -> collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
```

Compose sparse-selection replay with existing activation-checkpoint contexts.

**Parameters:**

**`context_fn`** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None`

Existing forward/recompute context factory, or None.

---

**Returns:** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None`

A factory owning a fresh selection recorder for every checkpointed

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

**`head_dim`**

---

**`indexer`**

---

**`k_norm`**

---

**`k_proj`**

---

**`num_heads`** `= config.num_attention_heads`

---

**`num_kv_heads`** `= config.num_key_value_heads`

---

**`o_proj`**

---

**`q_norm`**

---

**`q_proj`**

---

**`use_qk_norm`** `= getattr(config, 'use_qk_norm', False)`

---

**`v_proj`**

---

```python
nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Attention._initialize_attention() -> tuple[torch.nn.Module | None, typing.Any]
```

Return the `(attn_module, attn_func)` of `backend.attn`; a layer with its own kernels overrides this.

```python
nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Attention._project_qkv(
    x: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Project x\[..., hidden] to q\[..., num\_heads, head\_dim] and k/v\[..., num\_kv\_heads, head\_dim].

Per-head QK norm (over head\_dim) is applied before RoPE, matching the sglang reference
(`_qk_norm` then `rotary_emb`).

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

Map x\[batch, sequence, hidden] and freqs\_cis\[batch, sequence, rotary] to output\[batch, sequence, hidden].

**Parameters:**

**`x`** `torch.Tensor`

Hidden states; with `msa` they are packed to \[tokens, hidden] here and unpacked before
returning, so padding rows come back as zero.

---

**`freqs_cis`** `torch.Tensor`

Rotary table, packed alongside `x`.

---

**`attention_mask`** `torch.Tensor | None` — default: None

\[batch, sequence] or \[batch, heads, sequence, sequence]; None with `msa`,
whose documents travel as `cu_seqlens` in `attn_kwargs` instead.

---

**`msa`** `MSAMicrobatch | None` — default: None

The model-owned packed microbatch, given to every attention layer once MSA is on.

---

**`**attn_kwargs`** `Any` — default: \{}

Backend arguments.

---

```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
`build_block_sparse_attn_bias` to select which key blocks each query
attends to.

**`block_size`** `= sparse_cfg['sparse_block_size']`

---

**`index_head_dim`** `= sparse_cfg['sparse_index_dim']`

---

**`index_k_norm`**

---

**`index_k_proj`**

---

**`index_q_norm`**

---

**`index_q_proj`**

---

**`init_blocks`** `= sparse_cfg.get('sparse_init_block', 0)`

---

**`local_blocks`** `= sparse_cfg.get('sparse_local_block', 1)`

---

**`num_index_heads`** `= sparse_cfg['sparse_num_index_heads']`

---

**`score_type`** `= sparse_cfg.get('sparse_score_type', 'max')`

---

**`topk_blocks`** `= sparse_cfg['sparse_topk_blocks']`

---

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

Map x\[batch, sequence, hidden] and freqs\_cis\[batch, sequence, rotary] to a bool keep-mask
\[batch, num\_q\_heads, sequence, sequence].

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

```python
nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3Indexer.project_qk(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    cp_size: int,
    cp_rank: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Project x\[..., hidden] and freqs\_cis\[..., rotary] to rotated index\_q\[..., num\_index\_heads, index\_head\_dim]
and index\_k\[..., 1, index\_head\_dim]; `...` is \[batch, sequence] or \[tokens].

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

**`alpha`** `= float(getattr(config, 'swiglu_alpha', 1.702))`

---

**`down_proj`**

---

**`gate_proj`**

---

**`limit`** `= float(getattr(config, 'swiglu_limit', 7.0))`

---

**`up_proj`**

---

```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.MiniMaxM3MSAAttention(
    config: typing.Any,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [MiniMaxM3Attention](#nemo_automodel-components-models-minimax_m3_vl-layers-MiniMaxM3Attention)

A sparse attention layer running MSA's SM100 kernels on compact packed documents.

Same interface as `MiniMaxM3Attention`. Rejections are layered by when the answer can change:
backend and topology at construction, context parallelism at `setup_cp_attention`, and runtime
arguments once per microbatch in `MSAMicrobatch.build`.

```python
nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3MSAAttention._initialize_attention() -> tuple[None, None]
```

Own no generic backend: building one would import TransformerEngine before any MSA rejection
can fire, and a leftover `DotProductAttention` would send `apply_cp` down TE's context-parallel
branch (moe/parallelizer.py) instead of the `setup_cp_attention` rejection below.

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

Map x\[batch, sequence, hidden] and freqs\_cis\[batch, sequence, rotary] to output\[batch, sequence, hidden].

Both are packed to \[tokens, ...] by `msa` and the output unpacked with zero padding rows.
`attention_mask` and `attn_kwargs` are ignored: document isolation travels in `msa`.

```python
nemo_automodel.components.models.minimax_m3_vl.layers.MiniMaxM3MSAAttention.setup_cp_attention(
    cp_mesh: typing.Any
) -> None
```

Reject context parallelism; `apply_cp` dispatches on this method's presence and would otherwise
only log a warning, leaving CP silently unapplied.

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

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

**`_norm`**

---

**`weight`**

---

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

Normalize the last dimension with FP32 variance and weight arithmetic.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[..., hidden], with arbitrary leading dimensions.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., hidden], preserving the activation dtype.

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

```python
nemo_automodel.components.models.minimax_m3_vl.layers._compiled_rms_norm() -> collections.abc.Callable[[torch.Tensor, torch.Tensor, float, bool], torch.Tensor]
```

Share one lazily compiled stateless core while keeping eager instances eager.

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

```python
nemo_automodel.components.models.minimax_m3_vl.layers._rms_norm_core(
    x: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    gemma: bool
) -> torch.Tensor
```

Evaluate M3's FP32 normalization and scale before the final activation cast.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[..., hidden], with arbitrary leading dimensions.

---

**`weight`** `torch.Tensor`

Tensor of shape \[hidden], zero-centered when gemma is true.

---

**`eps`** `float`

Variance epsilon.

---

**`gemma`** `bool`

Whether to add one to the FP32 scale.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., hidden], with x's dtype and independently owned storage.

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

**`idx_q`** `torch.Tensor`

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

---

**`idx_k`** `torch.Tensor`

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

---

**`num_q_heads`** `int`

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

**`idx_q`** `torch.Tensor`

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

---

**`idx_k`** `torch.Tensor`

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

---

**`q_positions`** `torch.Tensor | None` — default: None

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