> 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.msa 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 | Name | Description | | ------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------ | | [`MSAMicrobatch`](#nemo_automodel-components-models-minimax_m3_vl-msa-MSAMicrobatch) | One packed microbatch's MSA state, built once and shared by every attention layer and stage. | | [`_SelectionPlan`](#nemo_automodel-components-models-minimax_m3_vl-msa-_SelectionPlan) | The layer-invariant half of block selection: the FMHA plan, the score shape and each query's geometry. | | [`_SparseAttention`](#nemo_automodel-components-models-minimax_m3_vl-msa-_SparseAttention) | The official flat forward, its saved schedule, and the backward-only aligned K/V workspace. | ### Functions | Name | Description | | ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | | [`_aligned`](#nemo_automodel-components-models-minimax_m3_vl-msa-_aligned) | Scatter `compact[tokens, H, D]` to rows `positions[tokens]` of a zero-filled `[workspace_size, H, D]`. | | [`_block_causal_documents`](#nemo_automodel-components-models-minimax_m3_vl-msa-_block_causal_documents) | Decode the document map of the dense mask a packed loader builds for a single-document pack. | | [`_contiguous_runs`](#nemo_automodel-components-models-minimax_m3_vl-msa-_contiguous_runs) | Check that every document is one contiguous run of tokens. | | [`_document_map`](#nemo_automodel-components-models-minimax_m3_vl-msa-_document_map) | Recover the int64 canonical document map `[batch, sequence]` (0 = padding) of one microbatch. | | [`_reject_unsupported_runtime`](#nemo_automodel-components-models-minimax_m3_vl-msa-_reject_unsupported_runtime) | Reject cache/THD/window/cross-attention/capture; tensor kwargs are checked only for presence. | | [`_score_scratch`](#nemo_automodel-components-models-minimax_m3_vl-msa-_score_scratch) | Return the process-wide score buffer of `device`, grown to `shape`. | | [`_warm_scorer`](#nemo_automodel-components-models-minimax_m3_vl-msa-_warm_scorer) | Compile every reachable scorer variant once per process and device, \~44 s each on a cold cache. | | [`require_msa_support`](#nemo_automodel-components-models-minimax_m3_vl-msa-require_msa_support) | Reject, at construction, an attention layer or backend the MSA kernels are not built for. | | [`sparse_attention`](#nemo_automodel-components-models-minimax_m3_vl-msa-sparse_attention) | Run MSA sparse attention on compact tokens. | ### Data [`_CACHE_ARGUMENTS`](#nemo_automodel-components-models-minimax_m3_vl-msa-_CACHE_ARGUMENTS) [`_CROSS_ATTENTION_ARGUMENTS`](#nemo_automodel-components-models-minimax_m3_vl-msa-_CROSS_ATTENTION_ARGUMENTS) [`_MEMO`](#nemo_automodel-components-models-minimax_m3_vl-msa-_MEMO) [`_SCORE_SCRATCH`](#nemo_automodel-components-models-minimax_m3_vl-msa-_SCORE_SCRATCH) [`_WARMUP_MAX_DOCS`](#nemo_automodel-components-models-minimax_m3_vl-msa-_WARMUP_MAX_DOCS) ### API ```python 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` --- ```python 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. ```python nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch.from_document_map( doc_ids: torch.Tensor, forced_blocks: tuple[int, int] ) -> nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch ``` 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. ```python 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. ```python 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. ```python 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. ```python 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]` --- ```python nemo_automodel.components.models.minimax_m3_vl.msa._SelectionPlan.build( msa: nemo_automodel.components.models.minimax_m3_vl.msa.MSAMicrobatch ) -> nemo_automodel.components.models.minimax_m3_vl.msa._SelectionPlan ``` classmethod Plan the scorer for `msa` and derive its `[tokens, blocks]` candidate and forced masks. ```python 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. ```python 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. ```python 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. ```python 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]. ```python 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]`. ```python 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. ```python 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 ```python 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. ```python 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. ```python 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. ```python 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. ```python 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. ```python 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. ```python nemo_automodel.components.models.minimax_m3_vl.msa._CACHE_ARGUMENTS = ('past_key_values', 'cache_position', 'page_table', 'seqused_k', 'prefix_cache') ``` ```python nemo_automodel.components.models.minimax_m3_vl.msa._CROSS_ATTENTION_ARGUMENTS = ('encoder_hidden_states', 'key_value_states') ``` ```python nemo_automodel.components.models.minimax_m3_vl.msa._MEMO: WeakIdKeyDictionary = WeakIdKeyDictionary() ``` ```python nemo_automodel.components.models.minimax_m3_vl.msa._SCORE_SCRATCH: dict[device, Tensor] = {} ``` ```python nemo_automodel.components.models.minimax_m3_vl.msa._WARMUP_MAX_DOCS = (16, 32, 64, 128, 256) ```