IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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).

Both ops are dual-arch, 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 attention op forks here, on has_amd_gpu_accelerator(); its AMD prefill plan is the arch-neutral msa_sm100_prefill_plan (host sizing + buffer alloc only -- no SM100 device kernel), shared by both arches, so only the pure-device run differs. The indexer op forks at both levels: USE_AMD_MTP_SCORER picks the AMD-only speculative route here, because that route is a different pair of kernels rather than a different lowering of the same one; every other indexer route forks one level down, inside sparse_indexer_{prefill,decode} and their _mtp siblings, which dispatch to the *_amd scorers and bitonic top-k on gfx950.

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-8 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 = 8):

  • == 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 through 8 -> 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-8 is correctly handled by this path, so no prefill/spec disambiguation is needed.
  • > 8 -> PREFILL (msa_sm100_prefill_{plan,run}, device CSR).

comptime values​

MAX_SPEC_DRAFT​

comptime MAX_SPEC_DRAFT = 8

Structs​

Models

Cookbook

Puzzles

Blog

Community

Contact