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

Context-parallel support for MiniMax M3 block-sparse DSA attention.

Under context parallelism the sequence is sharded across CP ranks with a
*load-balanced* layout (PyTorch's causal CP splits the sequence into
`2 * cp_size` chunks and assigns rank `r` the pair `{r, 2*cp_size-1-r}`),
so a rank's local positions are **not** a contiguous global span. The M3
lightning indexer builds its block-sparse mask from index q/k over the *global*
causal sequence, so a CP-aware sparse layer must gather the indexer inputs from
every rank and reorder them into global token order before selecting blocks.

This module holds the reorder primitives shared by the CP-aware attention. The
reorder math (`order_by_positions` / `restore_by_positions`) is factored out
as pure tensor functions so the load-balance inverse -- a silent-failure trap: a
wrong inverse trains without shape errors but never converges -- is unit-testable
on CPU without a process group.

## Module Contents

### Classes

| Name                                                                                                                      | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`MiniMaxM3CPSparseAttention`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-MiniMaxM3CPSparseAttention) | Context-parallel-aware drop-in for a MiniMax M3 sparse-attention layer. |
| [`_AllGatherConcatFn`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_AllGatherConcatFn)                 | All-gather + concat with an autograd-safe backward.                     |

### Functions

| Name                                                                                                                            | Description                                                                                          |
| ------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------- |
| [`_all_gather_concat_nograd`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_all_gather_concat_nograd)         | Plain (non-differentiable) all-gather + concat along `dim`.                                          |
| [`_block_sel_flat_offset`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_block_sel_flat_offset)               | Flat offset into a `[B, num_index_heads, query_length, num_blocks]` block-selection table, in int64. |
| [`_get_compiled_flash_attention`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_get_compiled_flash_attention) | Return static-shape FlexAttention compiled for the optional FLASH backend.                           |
| [`_get_compiled_flex_attention`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_get_compiled_flex_attention)   | -                                                                                                    |
| [`cp_document_ids`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-cp_document_ids)                             | Per-token document id from packed position ids (reset to 0 per document).                            |
| [`cp_load_balanced_global_slots`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-cp_load_balanced_global_slots) | Global token-slot indices for PyTorch's causal context-parallel load balancing.                      |
| [`order_by_positions`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-order_by_positions)                       | Reorder a CP-gathered tensor from load-balanced order into global token order.                       |
| [`restore_by_positions`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-restore_by_positions)                   | Select rows of a global-ordered tensor back into an arbitrary (local) position order.                |

### Data

[`_COMPILED_FLASH_ATTENTION`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_COMPILED_FLASH_ATTENTION)

[`_COMPILED_FLEX_ATTENTION`](#nemo_automodel-components-models-minimax_m3_vl-cp_sparse_attn-_COMPILED_FLEX_ATTENTION)

### API

```python
class nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention(
    config: typing.Any,
    backend: typing.Any,
    is_sparse_attention_layer: bool = True
)
```

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

Context-parallel-aware drop-in for a MiniMax M3 sparse-attention layer.

Inherits every parameter from `MiniMaxM3Attention`. Supported local CUDA
SDPA/FA4 cases use FlexAttention without a quadratic per-head keep-mask.
Other SDPA cases retain the eager forward; unsupported local FA4 cases reject.
`_cp_mesh` is installed post-FSDP via
`setup_cp_attention` (called by the MoE parallelizer's `apply_cp`).

Under CP (`cp_size > 1`) the sequence is sharded across ranks, so the DSA
block selection -- which is causal over the *global* sequence -- cannot be
built from a rank's local shard. This forward instead:

1. projects q/k/v + indexer q/k locally and applies QK-norm + RoPE locally
   (`freqs_cis` already encodes each token's global position, so phases
   stay correct after gathering);
2. all-gathers k/v (autograd-safe) and the indexer key + token positions
   across the CP group, then reorders them into global token order
   (load-balanced CP sharding is non-contiguous -- see
   `order_by_positions`);
3. selects the top-k key blocks for the *local* queries against the global
   key sequence (`select_sparse_blocks`);
4. attends with FlexAttention over a `BlockMask` that encodes the block
   selection + token-level causal, with the local queries against the full
   gathered K/V (`enable_gqa=True`). FlexAttention has a real backward,
   so the gathered K/V gradients flow back to the local shards.

Dense layers (0-2) are untouched; they use the standard DTensor-SDPA CP path.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention._cp_forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
```

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention._flex_sparse_attention(
    q: torch.Tensor,
    k_global: torch.Tensor,
    v_global: torch.Tensor,
    block_sel: torch.Tensor,
    q_positions: torch.Tensor,
    key_valid: torch.Tensor | None = None,
    doc_global: torch.Tensor | None = None,
    q_doc: torch.Tensor | None = None,
    keep_mask: torch.Tensor | None = None,
    rescue_pad_queries: bool = True
) -> torch.Tensor
```

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention._local_sparse_forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None
) -> torch.Tensor
```

Run local BSHD sparse attention without a `[batch, heads, sequence, sequence]` mask.

Block selection is unchanged. The optional caller keep-mask is composed inside
FlexAttention, including packed document masks and completely masked padding rows.

**Parameters:**

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

Hidden states \[batch, sequence, hidden].

---

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

Non-fused rotary cos/sin table \[batch, sequence, rotary\_dim].

---

**`attention_mask`** `torch.Tensor | None`

Optional \[batch, sequence] padding keep-mask or boolean
mask broadcastable to \[batch, heads, sequence, sequence].

---

**Returns:** `torch.Tensor`

Attention output \[batch, sequence, hidden], with zero output for fully

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention._select_local_blocks(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    q_positions: torch.Tensor
) -> torch.Tensor
```

Select key blocks once per checkpointed local attention call.

Only the final non-differentiable selection is retained for recompute.
Excluding the selection-only branch from TorchDispatch checkpoint modes
prevents selective checkpointing from saving its quadratic score matrix.
Trainable main Q/K/V projections and attention stay outside this region.

**Parameters:**

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

Hidden states \[batch, sequence, hidden].

---

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

Rotary cosine/sine table \[batch, sequence, rotary\_dim].

---

**`q_positions`** `torch.Tensor`

Global token slots \[sequence] for the local sequence.

---

**Returns:** `torch.Tensor`

Boolean selected blocks \[batch, index\_heads, sequence, key\_blocks],

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention.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.cp_sparse_attn.MiniMaxM3CPSparseAttention.setup_cp_attention(
    cp_mesh: typing.Any
) -> None
```

Install the CP submesh consumed by `_cp_forward` (model-owned CP).

Called post-FSDP by the MoE parallelizer's `apply_cp` for each sparse
layer. Routing M3 through this hook -- rather than having `apply_cp` set
`_cp_mesh` directly -- keeps it on the same model-owned CP path as the
other custom-attention models (Gemma4, DeepSeek-V4).

```python
class nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._AllGatherConcatFn()
```

**Bases:** `Function`

All-gather + concat with an autograd-safe backward.

Forward concatenates equal-sized local shards from all CP ranks along
`dim`. Backward all-reduces the concatenated gradient and slices out this
rank's shard. Mirrors `qwen3_5_moe/cp_linear_attn.py`'s helper.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._AllGatherConcatFn.backward(
    ctx,
    grad_output: torch.Tensor
)
```

staticmethod

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._AllGatherConcatFn.forward(
    ctx,
    local_tensor: torch.Tensor,
    group: 'dist.ProcessGroup',
    dim: int
)
```

staticmethod

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._all_gather_concat_nograd(
    tensor: torch.Tensor,
    group: 'dist.ProcessGroup',
    dim: int
) -> torch.Tensor
```

Plain (non-differentiable) all-gather + concat along `dim`.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._block_sel_flat_offset(
    batch_idx: torch.Tensor,
    index_head: torch.Tensor,
    query_idx: torch.Tensor,
    kv_block: torch.Tensor,
    num_index_heads: int,
    query_length: int,
    num_blocks: int
) -> torch.Tensor
```

Flat offset into a `[B, num_index_heads, query_length, num_blocks]` block-selection table, in int64.

FlexAttention inlines `mask_mod` into its Triton template and emits the inlined
index arithmetic in int32. A 4-D gather such as `block_sel[b, h, q, kb]` is
lowered to `((b * H + h) * S + q) * NB + kb`, which wraps once the table has more
than `2**31` elements: with four index heads and 128-token blocks that is a single
sequence of 262,144 tokens. Just past the limit the wrapped address still lands in
mapped memory, so tail queries silently read a wrong selection before the failure
escalates to an illegal memory access. Widening the operands here keeps the
generated address arithmetic in int64 (same fix as the Qwen3.8 flex QSA mask).

**Parameters:**

**`batch_idx`** `torch.Tensor`

Scalar batch coordinate of the table (after any head grouping).

---

**`index_head`** `torch.Tensor`

Scalar index-head coordinate.

---

**`query_idx`** `torch.Tensor`

Scalar local query coordinate, already clamped in range.

---

**`kv_block`** `torch.Tensor`

Scalar key block coordinate (`kv_idx // block_size`), already clamped.

---

**`num_index_heads`** `int`

Size of the table's index-head axis.

---

**`query_length`** `int`

Size of the table's query axis.

---

**`num_blocks`** `int`

Size of the table's key-block axis.

---

**Returns:** `torch.Tensor`

int64 offset of `[batch_idx, index_head, query_idx, kv_block]` in the flattened table.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._get_compiled_flash_attention() -> collections.abc.Callable[..., torch.Tensor]
```

Return static-shape FlexAttention compiled for the optional FLASH backend.

FLASH currently rejects symbolic scalar captures in the sparse mask. Keep
its shape-specialized cache separate from the existing dynamic Triton/CP
cache; changing sequence length explicitly compiles another specialization.

**Returns:** `Callable[..., torch.Tensor]`

Compiled FlexAttention callable for local sparse CUDA attention.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._get_compiled_flex_attention()
```

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.cp_document_ids(
    positions: torch.Tensor
) -> torch.Tensor
```

Per-token document id from packed position ids (reset to 0 per document).

`doc_id = cumsum(positions == 0) - 1` along the sequence dim: a 0-based id
that increments at every position-0 (document start). A single sequence -> all
zeros (so a same-document mask is all-True, a no-op). A trailing cp-pad (also
position 0) opens a spurious extra document, but pad keys/queries are excluded
by causality / the padding mask, so it is harmless.

**Parameters:**

**`positions`** `torch.Tensor`

`[B, T]` long global-ordered position ids.

---

**Returns:** `torch.Tensor`

`[B, T]` long document ids.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.cp_load_balanced_global_slots(
    cp_size: int,
    t_local: int,
    device: torch.device,
    rank: int | None = None
) -> torch.Tensor
```

Global token-slot indices for PyTorch's causal context-parallel load balancing.

Causal CP splits the (cp-padded) sequence into `2 * cp_size` equal chunks and
assigns rank `r` the pair `{r, 2*cp_size-1-r}` (concatenated in that order),
so the local length is `2 * chunk`. This reconstructs each local slot's global
index *structurally* -- independent of `position_ids` values -- which is robust
to cp-padding (pad slots land at the global tail, where causality excludes them)
and to the indexer's pad `position_id` fill.

**Parameters:**

**`cp_size`** `int`

context-parallel size.

---

**`t_local`** `int`

local (per-rank) sequence length; must be even.

---

**`rank`** `int | None` — default: None

if given, return the `[t_local]` slots for that CP rank; otherwise
return the `[cp_size * t_local]` slots for the rank-major all-gathered
concatenation (rank 0's tokens, then rank 1's, ...).

---

**Returns:** `torch.Tensor`

1-D long tensor of global slot indices (a permutation of `0..T_global-1`).

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.order_by_positions(
    gathered: torch.Tensor,
    gathered_positions: torch.Tensor,
    seq_dim: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Reorder a CP-gathered tensor from load-balanced order into global token order.

**Parameters:**

**`gathered`** `torch.Tensor`

tensor whose `seq_dim` concatenates every CP rank's local shard
in rank order (the output of an all-gather+concat).

---

**`gathered_positions`** `torch.Tensor`

1-D global token positions aligned with `gathered`
along `seq_dim` (gathered the same way). Must be a permutation of
`0..S-1` where `S = gathered.size(seq_dim)`.

---

**`seq_dim`** `int`

the sequence dimension of `gathered`.

---

**Returns:** `torch.Tensor`

`(global_tensor, sort_order)` where `global_tensor` is `gathered`

**Raises:**

* `ValueError`: if `gathered_positions` is not a dense permutation of 0..S-1.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.restore_by_positions(
    global_tensor: torch.Tensor,
    target_positions: torch.Tensor,
    seq_dim: int
) -> torch.Tensor
```

Select rows of a global-ordered tensor back into an arbitrary (local) position order.

Inverse companion to `order_by_positions`. Given a tensor indexed by
global position along `seq_dim` (position `p` at index `p`), return the
slice in `target_positions` order -- e.g. this rank's load-balanced local
positions, recovering the CP-sharded layout.

**Parameters:**

**`global_tensor`** `torch.Tensor`

tensor indexed by global position along `seq_dim`.

---

**`target_positions`** `torch.Tensor`

1-D positions to select, in the desired output order.

---

**`seq_dim`** `int`

the sequence dimension of `global_tensor`.

---

**Returns:** `torch.Tensor`

`global_tensor` gathered along `seq_dim` at `target_positions`.

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._COMPILED_FLASH_ATTENTION: Callable[..., Tensor] | None = None
```

```python
nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn._COMPILED_FLEX_ATTENTION = None
```