feat(minimax-m3): add SM100 MSA sparse attention with custom backward

[Butterfingrz]

Summary

This PR adds an opt-in SM100 training backend for MiniMax-M3 sparse-attention layers. With
BackendConfig.sparse_attn="msa", the attention operator uses the official MiniMax Sparse Attention (MSA)
prefill forward and an AutoModel-owned CuTe DSL backward with device-side task construction and gradient
finalization. This avoids expanding sparse block selection into a dense per-head token mask for the
sparse-SDPA path. The default remains sparse_attn="generic".

MSA applies to sparse layers; dense layers continue to use backend.attn. The model owns packed-document
layout and padding, while a private attention Adapter owns kernel dispatch and the state needed by backward.
The example recipe uses SDPA for dense layers and MSA for sparse layers, with 16k text packing and PP4/EP32.

Changelog

• Add the optional BackendConfig.sparse_attn selection and Linux msa dependency extra.
• Expose canonical document-local sparse block indices. The generic path converts the selection to its existing
  boolean mask representation; MSA consumes compact q2k indices.
• Add the model-private _msa.py module for packed-document layout, tensor/runtime validation, lazy dependency
  loading, and forward/backward state. Model and attention layers reuse its typed layout during each forward.
• Add the SM100 backward and its model-private launch sequence: delta preprocessing, task-table construction
  from the saved forward schedule, the main attention backward, and BF16 gradient finalization. The eligible
  fused task builder keeps exact counts and CTA scheduling on the device. dQ uses packed 16-bit atomic
  accumulation, while dK/dV retain FP32 accumulation.
• Add examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_tulu3_text_msa_16k.yaml: text-only Tulu-3 SFT,
  FSDP2 with PP4/EP32, CP1, and activation checkpointing. Its CI metadata requests 32 nodes with four processes
  per node on the GB200 cluster. Update the PyTorch container lock for the optional MSA dependency.

Design

MiniMaxM3TextModel resolves document identity and padding once per forward. Each pipeline stage that executes
MSA builds one _MSAPackedLayout and shares it with its sparse layers. Stages containing only dense layers
validate the document map without materializing the MSA layout. Callers supply ordinary model inputs and masks;
_msa_layout is model-owned.

Sparse layers compact hidden states and RoPE tables before projection, perform attention on real token rows,
and restore the attention output after o_proj. Hidden states enter and leave the attention layer as
[B, S, hidden]. The external attention format is qkv_format="bshd"; internally, Q/K/V use compact THD tensors,
where T is the number of real tokens:

MiniMaxM3TextModel.forward                           model.py
  document map -> _MSAPackedLayout                   _msa.py
        |
MiniMaxM3Attention.forward                           layers.py
  pack hidden states [B,S,hidden] -> [T,hidden]
  project Q/K/V; indexer selects document-local q2k
        |
_MSAFlatAttention(q, k, v, q2k, layout=layout)         _msa.py
  Q[T,64,128], K/V[T,4,128], int32 q2k[4,T,16]
  forward: official MSA CSR/schedule -> O[T,64,128]
        |
  o_proj -> unpack -> attention output [B,S,hidden]   layers.py

Autograd backward from saved compact state            _msa.py
  align attention K/V per document
  _run_msa_backward                                  kernels/msa_backward_sm100.py
    prepare per-call storage and compiled executables
    preprocess delta -> zero gradient pools -> build tasks
    main backward -> finalize BF16 gradients
  gather workspace dK/dV -> compact dQ/dK/dV

The flat-attention Interface consists of compact Q/K/V, document-local q2k[4,T,16], and the packed layout,
under the fixed topology and runtime constraints below.
q2k carries sparse support, with -1 in unused slots; the layout carries document offsets and token
coordinates. CSR, scheduler metadata, and backward task rows are derived execution metadata kept inside the
Adapter and its kernel Implementation. Layers do not construct or retain those execution structures.

_MSASparseAttentionFunction saves the compact tensors, forward output/LSE, and forward-derived CSR/schedule
for the corresponding backward pass. Backward creates attention K/V workspaces with each document padded to a
multiple of 128 tokens, computes gradients, and gathers dK/dV back into compact token order. The indexer may
separately align index-K during forward.

Supported scope

• SM100 (CC 10.0), BF16; 64 query / 4 KV / 4 index heads, head dimension 128, block size 128, top-k 16.
• Cache-free causal BSHD prefill; attention_dropout=0, rope_fusion=False, te_fp8=None, cp_size=1,
  num_mtp_modules=0. No sliding windows, CUDA graphs, or deterministic algorithms.
• Contiguous documents; multi-document packed rows crossing dense layers require backend.attn="sdpa"
  and a boolean [B,1,S,S] block-causal mask.
• Gradients use atomic accumulation: FP16 dQ by default (MSA_M3_DQ_ACCUM=bf16 selects wider-range BF16),
  FP32 dK/dV. Results are not bitwise deterministic.

Custom kernel

Model-local kernel

MiniMax-M3's custom SM100 backward lives under
nemo_automodel/components/models/minimax_m3_vl/kernels/.

Custom backward design

msa_backward_sm100.py owns the per-call storage and coordinates the kernels below. The autograd Adapter
supplies saved forward state and restores compact gradients after the kernel sequence.

_MSASparseAttentionFunction.backward                    _msa.py
  saved Q/K/V, O/LSE, forward CSR/schedule
  align attention K/V per document
        |
_run_msa_backward                                      msa_backward_sm100.py
  prepare temporary storage
  -> compute delta                                     msa_backward_preprocess_sm100.py
  -> zero gradient pools
  -> build task tables + CTA walk                      msa_task_build_sm100.py
       schedule contract / Torch fallback              msa_schedule.py
  -> main attention backward                           msa_backward_sm100.py
  -> de-interleave dQ; convert gradients to BF16        msa_grad_finalize_sm100.py
        |
  gather aligned dK/dV -> compact dQ/dK/dV               _msa.py

Dependency compatibility

The Linux msa extra pins official MSA 80434d7
and CuTe DSL 4.6.2, compatible with Quack 0.6.4. Dependencies load lazily;
kernels/msa_patch.py fixes upstream nvvm.fmax compatibility.

uv sync --locked --extra msa

Use uv run so build tools are on PATH.

TEST

Recorded local results; commands below are for reproduction. The logs do not pin kernel source or environment.

CPU unit tests:

CUDA_VISIBLE_DEVICES="" PYTHONPATH=. uv run --no-sync python -m pytest --cpu -q -ra \
  tests/unit_tests/models/minimax_m3_vl tests/unit_tests/moe/test_backend_config.py
# 197 passed, 4 skipped

Single-B200 SM100 functional tests:

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. uv run --no-sync python -m pytest -q -ra \
  tests/functional_tests/models/minimax_m3_vl/test_msa_sm100.py
# 3 passed, 1 skipped

The B200 tests cover packed O/dQ/dK/dV parity, top-16 selection, and checkpointed projection gradients.
PP2 was skipped because it requires two SM100 GPUs.

Additional generic CUDA checks: CP1 passed; test_eager_sparse_attn_bf16_matches_fp32 failed
(mean error 1.04e-3, threshold 1e-3), also reproduced at 6b17e7af.
The 128-GPU recipe was not run end to end.

Attention-operator benchmark

One B200, BF16, 64 Q / 4 KV heads, head_dim 128, block 128, top-16, batch 1, one document.
CUDA-event medians: 30 iterations after 5 warmups.

MSA uses fused task construction and FP16 dQ; generic uses dense-mask SDPA.
Both consume the same selection. Timing includes per-call preparation and excludes selection and projections.

T	MSA fwd / bwd ms	generic fwd / bwd ms	MSA speedup, fwd / bwd	peak GB, MSA / generic
4k	1.16 / 1.38	6.37 / 2.50	x5.5 / x1.8	1.3 / 3.7
8k	1.90 / 1.62	24.56 / 9.61	x12.9 / x5.9	2.6 / 13.9
16k	3.15 / 3.06	96.82 / 37.59	x30.7 / x12.3	5.2 / 53.5
32k	6.09 / 6.61	out of memory	—	10.5 / —

O/dQ/dK/dV relative L2 error: 1.9e-3–2.9e-3 at 4k–16k. Generic mask expansion takes 39.7 ms at 16k.

Peak allocation includes inputs and temporaries. At 32k, MSA uses 10.5 GB;
generic's 128 GiB bias allocation runs out of memory, preventing parity comparison.

Documentation and pre-checks

The model-local kernel README documents installation and compatibility. The MiniMax-M3 Fern model page still
needs user-facing MSA configuration and installation guidance in a follow-up.

• Make sure you read and followed Contributor guidelines.
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/claude review

[claude]

LGTM

Reviewed the full diff across all 17 non-uv.lock files. This PR adds the optional SM100 MSA sparse-attention backend for MiniMax M3 VL (CuTe DSL kernel + flat-prefill autograd adapter + packed-document layout + config/runtime gating + example recipe).

Key risk areas all check out:

• Backward determinism: FP32-atomic backward is nondeterministic and is explicitly rejected under torch.are_deterministic_algorithms_enabled() (tested), rather than silently shipped.
• Capabilities/parallelism: supports_pp/ep/cp already True; the new recipe (pp=4, ep=32, cp=1, tp=1, mtp=0) stays within supported axes, and MSA+MTP / MSA+CP are rejected with tests.
• Config-owned, typed boundaries: backend is copied to isolate the model from caller mutation; typed dataclasses cross the layout/schedule seams; the private _msa_layout is guarded from external callers.
• Optional dep: fmha_sm100 is gated via safe_import/safe_import_from, resolved lazily and cached, never at import time.
• Tests: forward+backward FP32 parity, truncated top-16 support, checkpointed layer param-grad parity, a real 2-rank PP subprocess parity gate, document-map precedence/rejection, and exact CTA row-cover partitioning.
• Ownership & contracts: all model-specific logic lives under components/models/minimax_m3_vl/; Python APIs carry Google-style tensor-layout docstrings, and the kernels (adapted from FlashAttention CuTe DSL) are densely annotated with inline shape/stride contracts.

No critical bugs, security issues, typed-API regressions, or coverage gaps found.

[HuiyingLi]

/ok to test 6b17e7a