nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn

Bases: 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.

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

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

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

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

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

nemo_automodel.components.models.minimax_m3_vl.cp_sparse_attn.MiniMaxM3CPSparseAttention.setup_cp_attention(
    cp_mesh: typing.Any
) -> None

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.

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

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

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

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:

Returns: torch.Tensor

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

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.

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:

Returns: torch.Tensor

[B, T] long document ids.

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:

Returns: torch.Tensor

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

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

Parameters:

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.

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:

Returns: torch.Tensor

global_tensor gathered along seq_dim at target_positions.