For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo module

msa

Graph-op bindings for the MiniMax-M3 block-sparse attention (MSA) kernels.

BF16, head_dim 128, paged KV with page_size 128 (== block size BN), single index-K head. Two ops:

mo.msa.indexer.ragged.paged -> sparse_indexer_{prefill,decode}
mo.msa.attention.ragged.paged -> msa_sm100_decode (decode) / msa_sm100_prefill_{plan,run} (prefill) on NVIDIA SM100; msa_amd_decode_dispatch (decode) / msa_amd_prefill_run (prefill) on AMD gfx950 (MI355).

The attention op is dual-arch: its body comptime-branches on has_amd_gpu_accelerator() so the dead arch's kernels never codegen (the SM100 tcgen05/TMA prefill cannot lower on gfx950, and the AMD MFMA path cannot lower on SM100). The AMD prefill plan is the arch-neutral msa_sm100_prefill_plan (host sizing + buffer alloc only -- no SM100 device kernel), shared by both arches; only the pure-device run differs. The indexer op is still SM100-only (the MSA indexer was not ported to AMD); an AMD M3 deployment must source d_indices from an AMD indexer or a host build -- the attention op only consumes d_indices.

Each op takes the same arguments for prefill and decode and picks the kernel at runtime from kv_collection.max_seq_length (the max number of new query tokens in the batch): == 1 is a single-token decode step, 2-4 is speculative decode, and anything larger is a prefill / context-encoding step.

The indexer op emits top-k block ids per (index head, token); the attention op consumes those block ids (d_indices) to gather a sparse band of KV blocks from the main paged cache. The index-K cache is always BF16; the main-KV cache is BF16 or native FP8 e4m3 (the attention op infers its kv_type from the operands). Neither carries scales, so both build with generic_get_paged_cache. FP8 here is scale-free and the msa_sm100_* kernels accumulate in FP32.

Modeled on the MLA FP8 indexer registration in attention.mojo (mo.mla.indexer.ragged.float8.paged) for the comptime cache-param extraction

paged-collection build, and on the in-tree MSA tests (Kernels/test/msa/test_msa_sm100_d128_decode_paged.mojo, test_msa_d128_prefill_device_csr.mojo) for the exact call shapes.

Three attention routes, picked at runtime from kv_collection.max_seq_length (max_q_len, the max number of new query tokens; MAX draft length = 4):

== 1 -> single-token DECODE (msa_sm100_decode, NullMask, the SM-fill split-K heuristic). Causal is a no-op (the single query sits at the sequence END, so every selected past KV position is causal-valid and nothing is masked).
2 / 3 / 4 -> sparse SPECULATIVE decode (msa_sm100_decode on NVIDIA, msa_amd_decode_dispatch on AMD) with spec_max_seq_len bound to the matched length: one CTA per (draft token, split-K partition), in-kernel per-token causal, capture-stable over-launched grid (batch * spec_max_seq_len on the token axis, max_num_partitions on the partition axis). At np>1 the partials key on the ragged global query row and the shared mha_splitk_reduce combine writes the ragged output directly. Causal is REAL (a draft token can precede some selected KV); the kernel derives each slot's logical KV start in-kernel as d_idx_base[blk] * BN and each token's logical query position as cache_lengths[batch] + tok_in_seq, so the op never builds a kv_logical_pos or q_positions array (mirrors the prefill use_causal path, which derives the diagonal from cu_seqlens + cache_lengths). A short prefill of 2-4 is correctly handled by this path, so no prefill/spec disambiguation is needed.
> 4 -> PREFILL (msa_sm100_prefill_{plan,run}, device CSR).

Structs

Struct_msa_attention_ragged_paged: Registers the mo.msa.attention.ragged.paged graph op with the graph compiler.
Struct_msa_indexer_ragged_paged: Registers the mo.msa.indexer.ragged.paged graph op with the graph compiler.

Structs​

Structs