ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMinimax M3 Vl

nemo_automodel.components.models.minimax_m3_vl.msa

View as Markdown

Packed-document MSA sparse attention for MiniMax M3 on SM100.

One MSAMicrobatch per microbatch holds everything MSA shares across the attention layers and the pipeline virtual stages: the canonical document map, the packed document layout, the lazily planned block scorer and the padding mask (ADR 0010). sparse_attention runs the official flat forward and the model-private backward on compact tokens; require_msa_support is the construction-time gate. Everything here is BSHD in and BSHD out; the compact [tokens, ...] layout lives between pack and unpack.

Module Contents

Classes

NameDescription
MSAMicrobatchOne packed microbatch’s MSA state, built once and shared by every attention layer and stage.
_SelectionPlanThe layer-invariant half of block selection: the FMHA plan, the score shape and each query’s geometry.
_SparseAttentionThe official flat forward, its saved schedule, and the backward-only aligned K/V workspace.

Functions

NameDescription
_alignedScatter compact[tokens, H, D] to rows positions[tokens] of a zero-filled [workspace_size, H, D].
_block_causal_documentsDecode the document map of the dense mask a packed loader builds for a single-document pack.
_contiguous_runsCheck that every document is one contiguous run of tokens.
_document_mapRecover the int64 canonical document map [batch, sequence] (0 = padding) of one microbatch.
_reject_unsupported_runtimeReject cache/THD/window/cross-attention/capture; tensor kwargs are checked only for presence.
_score_scratchReturn the process-wide score buffer of device, grown to shape.
_warm_scorerCompile every reachable scorer variant once per process and device, ~44 s each on a cold cache.
require_msa_supportReject, at construction, an attention layer or backend the MSA kernels are not built for.
sparse_attentionRun MSA sparse attention on compact tokens.

Data

_CACHE_ARGUMENTS

_CROSS_ATTENTION_ARGUMENTS

_MEMO

_SCORE_SCRATCH

_WARMUP_MAX_DOCS

API

class nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch(
    padding_mask: torch.Tensor,
    token_rows: torch.Tensor,
    workspace_positions: torch.Tensor,
    document_positions: torch.Tensor,
    document_workspace_starts: torch.Tensor,
    cu_seqlens: torch.Tensor,
    workspace_size: int,
    max_seqlen: int,
    forced_blocks: tuple[int, int]
)
Dataclass

One packed microbatch’s MSA state, built once and shared by every attention layer and stage.

_plan
_SelectionPlan

The scorer plan of this microbatch, built on the first block selection.

The first plan of a process also compiles every scorer variant production can reach: the from_pretrained path skips initialize_weights, so this is the one model-owned point both load paths pass through before the first scoring pass (ADR 0010).

cu_seqlens
Tensor
document_positions
Tensor
document_workspace_starts
Tensor
forced_blocks
tuple[int, int]
max_seqlen
int
padding_mask
Tensor
token_rows
Tensor
workspace_positions
Tensor
workspace_size
int
nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch.build(
    hidden: torch.Tensor,
    packed_seq_ids: torch.Tensor | None,
    attention_mask: torch.Tensor | None,
    padding_mask: torch.Tensor | None,
    attn_kwargs: collections.abc.Mapping[str, typing.Any],
    forced_blocks: tuple[int, int]
) -> nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch
classmethod

Return this microbatch’s state, building it once per batch tensor.

Every virtual pipeline stage is a deep-copied model that receives the same _packed_seq_ids tensor, so its identity plus _version keys the memo; the entry lives as long as the batch, which is why the state keeps no reference to the tensor. Without _packed_seq_ids the state is rebuilt on every call.

Parameters:

hidden
torch.Tensor

[batch, sequence, hidden] hidden states; only the shape and device are read.

packed_seq_ids
torch.Tensor | None

The loader’s [batch, sequence] document map, or None.

attention_mask
torch.Tensor | None

A 2-D document map or bool mask, a bool 4-D block-causal mask, or None.

padding_mask
torch.Tensor | None

bool [batch, sequence] padding mask, or None.

attn_kwargs
Mapping[str, Any]

The forward’s backend keyword arguments, checked for unsupported runtime features.

forced_blocks
tuple[int, int]

(init_blocks, local_blocks) of the model’s indexer.

Raises:

  • NotImplementedError: For non-BSHD input, caches, non-causal or windowed attention, or CUDA graph capture.
  • ValueError: If no source yields a well-formed document map.
classmethod

Derive the packed layout of doc_ids with exactly one device-to-host synchronization.

Parameters:

doc_ids
torch.Tensor

Integer [batch, sequence] document map: 0 marks padding, each positive id one contiguous run of tokens within its row.

forced_blocks
tuple[int, int]

(init_blocks, local_blocks) of the model’s indexer.

Raises:

  • ValueError: If the map is empty, holds negative ids, no real token, an interrupted document, or coordinates past int32.
nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch.pack(
    external: torch.Tensor
) -> torch.Tensor

Gather external[batch, sequence, ...] to [tokens, ...] in document order; may alias the input.

nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch.select_blocks(
    index_q: torch.Tensor,
    index_k: torch.Tensor
) -> torch.Tensor

Choose each query’s key blocks within its own document for one layer.

Selection is a hard top-k over unnormalized QK maxima, so it is not differentiable: call it under torch.no_grad.

Parameters:

index_q
torch.Tensor

bf16 [tokens, 4, 128] index queries, post norm and RoPE.

index_k
torch.Tensor

bf16 [tokens, 1, 128] shared index key, post norm and RoPE.

Returns: torch.Tensor

int32 [4, tokens, 16] document-local block ids, padded with -1: the canonical support.

nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch.unpack(
    packed: torch.Tensor
) -> torch.Tensor

Scatter packed[tokens, ...] back to [batch, sequence, ...] with zero padding; may alias the input.

class nemo_automodel.components.models.minimax_m3_vl.msa._SelectionPlan(
    plan: typing.Any,
    score_shape: tuple[int, int, int],
    num_blocks: int,
    candidate: torch.Tensor,
    forced: torch.Tensor
)
Dataclass

The layer-invariant half of block selection: the FMHA plan, the score shape and each query’s geometry.

Pinned to split_prefill_decode=False and num_kv_splits=1: the first splits a batch whose first document is short into two sub-plans (2.9x on the score pass plus two host syncs per call), the second lets the planner pick a variant from an SM-count estimate.

candidate
Tensor
forced
Tensor
num_blocks
int
score_shape
tuple[int, int, int]
classmethod

Plan the scorer for msa and derive its [tokens, blocks] candidate and forced masks.

nemo_automodel.components.models.minimax_m3_vl.msa._SelectionPlan.select(
    index_q: torch.Tensor,
    index_k: torch.Tensor
) -> torch.Tensor

Score bf16 index_q[tokens, 4, 128] against index_k[tokens, 1, 128] -> int32 [4, tokens, 16] block ids.

class nemo_automodel.components.models.minimax_m3_vl.msa._SparseAttention()

Bases: Function

The official flat forward, its saved schedule, and the backward-only aligned K/V workspace.

nemo_automodel.components.models.minimax_m3_vl.msa._SparseAttention.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor | None
) -> tuple[typing.Any, ...]
staticmethod

Map bf16 grad_out[tokens, 64, 128] to dq[tokens, 64, 128], dk/dv[tokens, 4, 128] and two None slots.

nemo_automodel.components.models.minimax_m3_vl.msa._SparseAttention.forward(
    ctx: typing.Any,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q2k: torch.Tensor,
    msa: nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch
) -> torch.Tensor
staticmethod

Run bf16 q[tokens, 64, 128], k/v[tokens, 4, 128], int32 q2k[4, tokens, 16] -> bf16 out[tokens, 64, 128].

nemo_automodel.components.models.minimax_m3_vl.msa._aligned(
    compact: torch.Tensor,
    positions: torch.Tensor,
    workspace_size: int
) -> torch.Tensor

Scatter compact[tokens, H, D] to rows positions[tokens] of a zero-filled [workspace_size, H, D].

nemo_automodel.components.models.minimax_m3_vl.msa._block_causal_documents(
    mask: torch.Tensor,
    shape: tuple[int, int]
) -> torch.Tensor

Decode the document map of the dense mask a packed loader builds for a single-document pack.

Kept until PR #3831 lets consumes_packed_seq_ids request the compact map for every pack; delete this decoder once that merges.

Parameters:

mask
torch.Tensor

bool [batch, 1, sequence, sequence]; query row i keeps key j where true.

shape
tuple[int, int]

The expected (batch, sequence).

Returns: torch.Tensor

int64 [batch, sequence] document ids, 0 for padding.

Raises:

  • ValueError: If the mask is not standard block-causal (every real query keeps exactly the causal keys of its own contiguous document, padding rows all false), checked with one host synchronization.
nemo_automodel.components.models.minimax_m3_vl.msa._contiguous_runs(
    ids: torch.Tensor,
    batch_rows: torch.Tensor,
    is_real: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Check that every document is one contiguous run of tokens.

Parameters:

ids
torch.Tensor

int64 [batch * sequence] flat document map.

batch_rows
torch.Tensor

int64 [batch * sequence] batch row of each flat token.

is_real
torch.Tensor

bool [batch * sequence]; True where ids > 0.

Returns: torch.Tensor

(valid, first_bad_row) 0-d tensors: whether no run is interrupted, and the first flat row

nemo_automodel.components.models.minimax_m3_vl.msa._document_map(
    hidden: torch.Tensor,
    packed_seq_ids: torch.Tensor | None,
    attention_mask: torch.Tensor | None,
    padding_mask: torch.Tensor | None
) -> torch.Tensor

Recover the int64 canonical document map [batch, sequence] (0 = padding) of one microbatch.

Only the shape and device of hidden [batch, sequence, hidden] are read. Sources, in priority order: the packed loader’s _packed_seq_ids, a 2-D attention_mask holding document ids or a bool 4-D block-causal attention_mask (decoded by _block_causal_documents), a bool padding_mask (one document per row), else one document per row.

nemo_automodel.components.models.minimax_m3_vl.msa._reject_unsupported_runtime(
    attn_kwargs: collections.abc.Mapping[str, typing.Any]
) -> None

Reject cache/THD/window/cross-attention/capture; tensor kwargs are checked only for presence.

nemo_automodel.components.models.minimax_m3_vl.msa._score_scratch(
    device: torch.device,
    shape: tuple[int, int, int]
) -> torch.Tensor

Return the process-wide score buffer of device, grown to shape.

The scorer stores rather than accumulates and every tile the selection rule reads is written by the same call, so one buffer serves every layer and microbatch; scoring passes never overlap because MSA is single-stream and rejects CUDA-graph capture. max_k_tiles is rounded up to 128 tiles whatever the documents are, so per-plan buffers would cost 224 MiB where this one costs 11.28 MiB.

nemo_automodel.components.models.minimax_m3_vl.msa._warm_scorer(
    device: torch.device,
    forced_blocks: tuple[int, int]
) -> None

Compile every reachable scorer variant once per process and device, ~44 s each on a cold cache.

Warming runs the production path over synthetic one-document microbatches, so the variant compiled here is the variant production reaches by construction.

nemo_automodel.components.models.minimax_m3_vl.msa.require_msa_support(
    attention: typing.Any,
    backend: nemo_automodel.components.models.common.BackendConfig
) -> None

Reject, at construction, an attention layer or backend the MSA kernels are not built for.

Parameters:

attention
Any

The sparse attention layer: num_heads, num_kv_heads, head_dim and an indexer with num_index_heads, block_size, topk_blocks, index_head_dim and score_type.

backend
BackendConfig

The model’s backend selection.

Raises:

  • ValueError: If the topology is not the 64-query/4-KV-head, 128-channel, top-16 one, or the block score is not the max reduction.
  • NotImplementedError: If the backend asks for FP8 projections or fused RoPE.
nemo_automodel.components.models.minimax_m3_vl.msa.sparse_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q2k: torch.Tensor,
    msa: nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch
) -> torch.Tensor

Run MSA sparse attention on compact tokens.

Parameters:

q
torch.Tensor

bf16 [tokens, 64, 128] queries after RoPE.

k
torch.Tensor

bf16 [tokens, 4, 128] keys after RoPE.

v
torch.Tensor

bf16 [tokens, 4, 128] values.

q2k
torch.Tensor

int32 [4, tokens, 16] canonical support from msa.select_blocks.

msa
MSAMicrobatch

The microbatch the tokens were packed by.

Returns: torch.Tensor

bf16 [tokens, 64, 128] attention output; q, k and v receive gradients.

Raises:

  • NotImplementedError: Under torch.use_deterministic_algorithms: the backward accumulates dK/dV with FP32 atomics and dQ with packed bf16 atomics, so it is not bitwise deterministic.
  • ValueError: If q, k or v is not bf16.
nemo_automodel.components.models.minimax_m3_vl.msa._CACHE_ARGUMENTS = ('past_key_values', 'cache_position', 'page_table', 'seqused_k', 'prefix_cache')
 
nemo_automodel.components.models.minimax_m3_vl.msa._CROSS_ATTENTION_ARGUMENTS = ('encoder_hidden_states', 'key_value_states')
 
nemo_automodel.components.models.minimax_m3_vl.msa._MEMO: WeakIdKeyDictionary = WeakIdKeyDictionary()
 
nemo_automodel.components.models.minimax_m3_vl.msa._SCORE_SCRATCH: dict[device, Tensor] = {}
 
nemo_automodel.components.models.minimax_m3_vl.msa._WARMUP_MAX_DOCS = (16, 32, 64, 128, 256)