jburtoft/minimax-m3-msa-neuron-kernels

MiniMax-M3 MSA Block-Sparse GQA Prefill Kernels for AWS Neuron

A Trainium port of MiniMax's Multi-Scale Sparse Attention (MSA) prefill attention, implemented in NKI 0.5 for trn2. Handles the top-k-KV-block-selection attention pattern used by MiniMax-M3 and other Native-Sparse-Attention / DSA-family models.

Three kernels:

• block_sparse_gqa_prefill_kernel -- SPMD (LNC=2) prefill with normalized output. Main kernel.
• block_sparse_gqa_prefill_state_kernel -- SPMD prefill with (o, l, m) state output for context-parallel merge.
• block_sparse_attention_single_q_kernel -- single-Q-block MHA POC with broader shape coverage.

Version: v1.0.0

Performance (MiniMax-M3 shape: H_q=64, H_kv=4, D=128, block_size=128, topk=16, trn2.3xlarge LNC=2, BF16)

Measured on SDK 2.31 DLAMI 20260708 + PyTorch Native Beta 4 (2026-08-05), warm median of 5 runs.

Single-core (one logical NeuronCore, kernel[2] at LNC=2 uses 2 bankmate physical cores)

seq_len	Blocks kept	NKI kernel	ms per Q-block
8,192	25.0%	77.5 ms	1.21
16,384	12.5%	153.0 ms	1.20
32,768	6.25%	413.2 ms	1.61
65,536	3.12%	823.9 ms	1.61

Constant ~1.2-1.6 ms per Q-block regardless of seq_len. Dense torch.compile(backend='neuron') compile-fails at seq_len >= ~12K (StableHLO broadcast_in_dim limit; see Known Limitations).

Multi-core Q-parallel (all 4 logical NeuronCores, 1 process per core)

At LNC=2, trn2.3xlarge exposes 4 logical cores. NKI SPMD (kernel[N]) is capped at N=2 (one logical core, its 2 bankmate physical cores), but you can run 4 independent processes, each pinned to a different logical core via NEURON_RT_VISIBLE_CORES, each computing a disjoint 1/4 slice of the Q sequence against the full KV cache. Output is bit-identical to the single-core kernel. See examples/05_qparallel.py.

seq_len	Single-core	Q-parallel DP=4	Speedup	Efficiency
8,192	77.5 ms	21.6 ms	3.59x	90%
16,384	153.0 ms	40.0 ms	3.83x	96%
32,768	413.2 ms	106.4 ms	3.88x	97%
65,536	823.9 ms	208.6 ms	3.95x	99%

Scaling improves with seq_len: at 64K prefill, Q-parallel gives 99% of ideal 4x speedup (208.6 ms wall-clock, ~314K tokens/sec single-request).

Context parallel (state kernel + cp_merge)

Via block_sparse_gqa_prefill_state_kernel. CP=2 correctness verified at 8K (cos_sim = 0.9999940 vs single-core). CP extends the seq_len ceiling from ~4.2M (single-node) to ~8.4M tokens FP32 on trn2.3xlarge. Orthogonal to Q-parallel: Q-parallel replicates the KV cache per rank; CP shards it.

Reproduce via examples/03_benchmark.py, examples/05_qparallel.py, and examples/04_cp_merge.py.

Correctness

• BF16 vs FP32 numpy reference (skipping the first topk-1 Q blocks that have sentinel slots; see Known Limitation #2):
  • Shape A (S=1024, H_q=8, H_kv=2, topk=4): cos_sim = 0.9999753, max_abs_diff = 3.4e-3
  • Shape B MiniMax-M3 (S=8192, H_q=64, H_kv=4, topk=16): cos_sim = 0.9999754, max_abs_diff = 2.1e-3
• Q-parallel DP=4 vs single-core: max_abs_diff = 0.0 (bit-identical) at every tested seq_len.
• CP=2 merged output vs single-node (Shape B): cos_sim = 0.9999940 min row-wise. max_abs_diff on the tail is O(0.1) due to BF16 rounding-order drift across ranks, but direction is essentially identical.
• Reproduce via examples/02_parity.py, examples/04_cp_merge.py, and examples/05_qparallel.py --verify.

Direct usage

import torch
from kernels import get_kernel

# revision + trust_remote_code required by kernels >= 0.15
m3msa = get_kernel(
    "jburtoft/minimax-m3-msa-neuron-kernels",
    revision="v1.0.0",
    trust_remote_code=True,
)

# MiniMax-M3 shape, BF16, single-node prefill.
B, S, H_q, H_kv, D = 1, 8192, 64, 4, 128
q = torch.randn(B, S, H_q, D, dtype=torch.bfloat16, device="neuron")
k = torch.randn(B, S, H_kv, D, dtype=torch.bfloat16, device="neuron")
v = torch.randn(B, S, H_kv, D, dtype=torch.bfloat16, device="neuron")

# kv_indices: [B, num_q_blocks, TOPK] int32, -1 sentinels for unused slots.
# Produce this from your Lightning-Indexer (or any block-level attention router).
# See ops.build_block_kv_indices() to convert a block-level keep mask.
kv_indices = ...  # e.g. m3msa.build_block_kv_indices(keep_mask, topk=16)

out = m3msa.block_sparse_gqa_prefill(q, k, v, kv_indices)
# out: [B, S, H_q, D] bfloat16

Context-parallel usage

# Each rank owns half of the KV cache. Indices for KV blocks not owned by
# this rank must be set to the sentinel -1 in the local kv_indices tensor.
o_r, l_r, m_r = m3msa.block_sparse_gqa_prefill_state(
    q, k_shard, v_shard, kv_indices_local,
)

# all_gather (o, l, m) across CP ranks, then merge on host:
states = [(o_r, l_r, m_r) for each rank]   # e.g. via torch.distributed.all_gather
out = m3msa.cp_merge(states, output_dtype=torch.bfloat16)

See examples/04_cp_merge.py for a full, single-device simulation of CP=2 with parity to a single-node run.

Repository layout

build/torch-neuron/
├── __init__.py                                  <- public API re-exports
├── metadata.json                                <- HF kernels library metadata
├── constants.py                                 <- BS, D, H_Q, H_KV, TOPK, LNC_DEGREE, sentinel
├── ops.py                                       <- host helpers: mask + index builders,
│                                                   numpy FP32 reference for parity
└── nki_kernels/
    ├── __init__.py
    ├── block_sparse_gqa_prefill.py              <- main SPMD kernel (normalized out)
    ├── block_sparse_gqa_prefill_wrapper.py      <- eager wrapper
    ├── block_sparse_gqa_prefill_state.py        <- SPMD kernel with (o, l, m) state out
    ├── block_sparse_gqa_prefill_state_wrapper.py <- eager wrapper + cp_merge()
    └── block_sparse_attention_single_q.py       <- single-Q-block POC kernel

examples/
├── README.md
├── _loader.py               <- local-clone-or-HF-Hub loader
├── 01_smoke_test.py         <- verify kernel loads + one forward
├── 02_parity.py             <- BF16 kernel vs FP32 numpy reference
├── 03_benchmark.py          <- reproduce the single-core perf table
├── 04_cp_merge.py           <- CP=2 correctness vs single-node
├── 05_qparallel.py          <- Q-parallel driver (4 processes, 1 per logical core)
└── launch_qparallel.sh      <- launcher for 05_qparallel.py

Users load the package via HF's kernels library -- the internal file split is transparent.

Reference config

The kernels are compiled against the MiniMax-M3 reference shape (defined in constants.py):

Constant	Value	Note
BS	128	block size (tokens per K/V block, also Q tile size)
D	128	head dim
H_Q	64	query heads per layer
H_KV	4	KV heads per layer (GQA group = 16)
TOPK	16	top-K KV blocks per Q block (from Lightning-Indexer)
LNC_DEGREE	2	SPMD degree (LNC=2 on trn2.3xlarge)

Requirements:

• S_q == S_k (self-attention prefill only)
• S_q % (BS * LNC_DEGREE) == 0 (divisible by 256 by default)
• H_q % H_kv == 0 (GQA)

Other GQA shapes trigger a recompile. The kernel specializes on the shapes it sees; nothing is hard-coded to 64/4/128, but correctness has only been verified on the reference shape and on Shape A (H_q=8, H_kv=2, D=128, TOPK=4).

How it works

The kernel implements a flash-attention-style algorithm restricted to a caller-selected subset of KV blocks:

1. Indirect DMA gather: for each Q-block, gather the TOPK selected KV blocks via .ap(vector_offset=..., indirect_dim=0) with oob_mode.skip for -1 sentinel slots. This is the primitive that physically skips masked KV blocks -- not additive -inf mask theater.
2. Per-instance HBM staging: restage into nl.private_hbm (not nl.shared_hbm -- see Known Limitations #1). Layout swings the partition dim from TOPK to BS, which the matmul path requires.
3. Online softmax: running max + running sum + running unnormalized output. Standard flash-attention accumulator, kept in FP32 in SBUF.
4. Intra-block causal masking: applied to the diagonal (q_blk == k_blk) slot only. Callers should place the local block at kv_indices[..., 0] by convention (the Lightning-Indexer's default).
5. Fused exp + row-sum: nisa.activation(op=exp, reduce_op=nl.add, reduce_cmd=reset_reduce) -- one pass over the score tile instead of two.
6. SPMD sharding: kernel[LNC_DEGREE](args) launches LNC_DEGREE SPMD instances, each handling NUM_Q_TILE / LNC_DEGREE Q-blocks (default: 1 Q-tile per instance).
7. Normalize (or emit state): the main kernel divides o_acc / l and returns the normalized output. The state kernel skips the division and returns (o_unnormalized, l, m) for host-side CP merge.

Kernel details

block_sparse_gqa_prefill_kernel

The main kernel. SPMD (kernel[2] at LNC=2), returns normalized [NUM_Q_TILE, BS, H_q, D]. The eager wrapper block_sparse_gqa_prefill(q, k, v, kv_indices) loops over NUM_Q_TILE-sized Q-tile batches for you.

At 8K MiniMax-M3 shape (BF16): 77.6 ms median warm. Dense torch.compile(backend='neuron') baseline compile-fails at seq_len >= ~12K.

block_sparse_gqa_prefill_state_kernel

Identical math but emits FP32 (o_unnormalized, l, m) instead of the normalized output. Standard flash-attention state, composable across disjoint KV shards via the merge formula in cp_merge():

m_final = elementwise_max(m_r for r in ranks)
l_final = sum(l_r * exp(m_r - m_final) for r in ranks)
o_final = sum(o_r * exp(m_r - m_final) for r in ranks) / l_final

CP=2 on trn2.3xlarge extends the seq_len ceiling from ~4.2M (single-node) to ~8.4M tokens FP32.

block_sparse_attention_single_q_kernel

Single-Q-block MHA POC, no GQA, no SPMD. Retained because:

• Runs on the NKI 0.5 CPU simulator (useful for correctness bring-up without touching hardware).
• Broader shape coverage (any BS, D, TOPK combination that fits in SBUF).
• Simpler starting point for porting to other block-sparse layouts (2D NATTEN, DSA, custom sparsity patterns).

Not perf-tuned; use the SPMD kernel for real workloads.

ops.build_block_kv_indices / ops.build_causal_block_mask / ops.dense_gqa_reference

Host-side utilities:

• build_block_kv_indices(keep_mask, topk): convert a [B, num_q_blocks, num_kv_blocks] bool mask to the [B, num_q_blocks, topk] int32 layout the kernels consume, with -1 sentinel padding.
• build_causal_block_mask(BS): the [BS, BS] additive intra-block causal mask.
• dense_gqa_reference(q, k, v, kv_indices, ...): numpy FP32 dense-attention-with-block-sparse-mask reference. Used by 02_parity.py.

Known limitations

1. nl.shared_hbm is coherent within a logical core's bankmate pair, but not synchronized on concurrent writes. On trn2 each "logical NeuronCore" (as reported by neuron-ls at LNC=2) is composed of 2 physical NeuronCores that share the same HBM stack ("bankmates"). A nl.shared_hbm allocation lives in that shared HBM stack and is a single physical buffer visible to both physical cores. When you launch kernel[2] at LNC=2, the two SPMD instances run on those two bankmate cores concurrently -- they can both read and write the same allocation coherently, but there is no automatic write-race prevention. If both instances write to the same address at once (as our per-instance K/V staging tensor did in an earlier draft), the result is a data race. Fix: use nl.private_hbm for per-instance scratch (each physical core gets its own buffer). Return values must still be nl.shared_hbm. Verified: disjoint-slice writes to shared_hbm are coherent; overlapping-slice writes race and produce mixed data. Note: kernel[N] at LNC=2 caps at N=2 (a single logical core's bankmate pair). To use more physical cores you need framework-level data parallelism (multiple processes or torch data-parallel), not NKI SPMD.
2. Sentinel semantics are "zero-fill", not -inf. -1 entries in kv_indices cause oob_mode.skip on the indirect DMA, and the destination SBUF is pre-zeroed via memset(0). So a sentinel slot contributes Q @ 0 = 0 to the score row, which gets a small softmax weight (roughly 1/N_active) and contributes ~0 * V[slot_0] to the output. This is not equivalent to strict -inf masking -- it produces a small O(1/TOPK) output bias when the sentinel-slot count is non-trivial. In practice, indexer-produced kv_indices are almost always fully populated (Lightning-Indexer emits exactly topk_blocks real indices per Q block), so this only bites at the sequence boundary when fewer than topk_blocks local blocks exist. The dense_gqa_reference in ops.py uses strict -inf masking, so parity tests should fully populate kv_indices (see examples/02_parity.py).
3. Runtime softmax_scale is not plumbed through. The kernels hard-code 1/sqrt(D). The eager wrappers raise NotImplementedError if a non-default scale is passed rather than silently ignoring it. Easy fix if needed -- open an issue.
4. Dense torch.compile(backend='neuron') compile-shape ceiling at ~12K: the StableHLO broadcast_in_dim op requires statically-shaped results and rejects [H, S, S] score tensors above roughly S=12K on SDK 2.31 Beta 4. This is a limitation of the dense comparison baseline, not of this NKI kernel -- the NKI kernel never materializes [H, S, S] and runs cleanly to 64K+.
5. Decode-side is not implemented. This is a prefill-only kernel. Decode-side integration in NxDI requires framework-level changes to KVCacheManager output-aliasing that are outside this kernel's scope.
6. Correctness verified only on the reference GQA shape and Shape A (H_q=8, H_kv=2, D=128, TOPK=4). Other shapes should compile (shape specialization) but their numerical outputs have not been checked against a reference.

Items 1, 2, 3, and 5 are documented; item 4 is a filed internal ticket against the SDK 2.31 compiler.

Environment (verified working)

• Instance: trn2.3xlarge, ap-southeast-4 (Melbourne)
• DLAMI: Deep Learning AMI Neuron (Ubuntu 24.04) 20260708 (SDK 2.31)
• PyTorch Native: Beta 4 (SageMaker Training DLC, 2026-07-23)
• Versions:
  • torch-neuronx 2.11.3.0.1419+d09a7917
  • neuronx-cc 2.26.6360
  • NKI 0.5.0
• LNC mode: LNC=2 (default on trn2.3xlarge)
• Venv: ~/beta35-venv/bin/activate (via deploy-beta35.sh --variant beta4)

Attribution

This is a Trainium port of the Multi-Scale Sparse Attention (MSA) pattern used in:

• MiniMax-M1 / MiniMax-M3 technical reports (MiniMax AI)
• The reference GPU implementation at MiniMaxAI/msa on HF Kernels

The block-sparse top-k-KV-block-selection pattern is also related to:

• DeepSeek DSA (Dense Sparse Attention, DeepSeek V3)
• FlashAttention-2 (Dao) -- online-softmax state and CP merge formula
• NATTEN (SHI Labs) -- neighborhood-attention block-sparse patterns

License

Apache 2.0. See LICENSE.