Kernel Optimization For Minimax M3 On Nvidia Blackwell

Optimizing MiniMax M3 Sparse Attention on NVIDIA Blackwell

GLM 5.2 Fast is available! Opus-level intelligence at open-source rates. No contracts, pay per token. Start building.

Blog

Kernel Optimization For Minimax M3 On Nvidia Blackwell Optimizing MiniMax M3 Sparse Attention on NVIDIA Blackwell

PUBLISHED 7/10/2026

Table of Contents 1. The algorithm and the kernel design space 2. The KV-outer kernel 3. Optimizations 4. Differences from open-source MSA 5. Results Acknowledgements Appendix: the I/O cost model

Table of Contents

Summary: Attention drives most compute and memory costs in long-context inference. Sparse attention, as used in MiniMax M3 , reduces this by having each query attend only to the top 16 most relevant 128-token KV blocks. Implementing this efficiently is difficult, however; the sparse selection is data-dependent and varies by request, leading to irregular memory access patterns that often negate the theoretical speedup. The Fireworks AI Performance team built a Blackwell (SM100) kernel for M3 sparse attention that closes most of that gap. It uses a KV-stationary execution path — a research idea brought up by the MiniMax team in the Minimax Sparse Attention paper — that loads each selected KV block once in the outer loop and attends every query that chose it in the inner loop. With this schedule fixed, performance optimization focuses on two critical areas: reducing gathered-load/scattered-store memory traffic and improving load balancing across SMs. Our implementation addresses both challenges, achieving the following results on a single B200 (fp8): • Attention Kernel Efficiency: Reaches throughput of ~980 TFLOP/s at ~4.1 TB/s HBM bandwidth. This represents a 1.9–2.4× speedup over a query-stationary baseline (FlashInfer) and a ~1.6× improvement over MiniMax’s open-source MSA kernel. • Full Module Performance: Including index mapping and combination stages, the module delivers 1.18–1.43× gains over the baseline and 1.32–1.41× gains over open-source MSA.

Figure [above-dark]: Fireworks KV-outer E2E sparse attention module achieves a 1.18–1.43× speedup over a query-stationary baseline (FlashInfer) and a 1.32–1.41× speedup over MiniMax's open-source MSA kernels. Figure [above-light]: Fireworks KV-outer main attention kernel reaches throughput of ~980 TFLOP/s at ~4.1 TB/s HBM bandwidth, achieving 1.6× speedup over Minimax's open-source MSA kernel (compared to the baseline FlashAttention-4). 1. The algorithm and the kernel design space

M3 sparse attention

M3 attention is blockwise sparse on top of grouped-query attention (GQA): the KV sequence is partitioned into fixed 128-token blocks, a lightweight index branch scores those blocks and selects the topK (= 16) highest-scoring blocks for each GQA group, and a main branch then computes exact attention over only the selected blocks. Shapes: Hq = 64, Hkv = 4 (GQA factor Hq/Hkv = 16), D = 128, block = 128. The algorithm and the KV-stationary execution idea are from MiniMax ( paper , MSA kernels ). Q-outer vs KV-outer

A traditional dense FlashAttention-style kernel (e.g. FlashAttention4, aka FA4) loops Q-outer / K-inner: one threadblock owns a 128-row Q tile and streams all KV through it at a 128×128 MMA. (128 is the common block/tile size we use as the running example; other block sizes are valid.) As a result each KV load is shared across the tile’s 128 Q tokens, and the MMA runs at a large, efficient 128×128 tile. Sparsity, on the other hand, admits two kernel structures: • Q-outer (query-stationary). Keep Q-outer/K-inner; each query reads only its selected KV blocks. The M dimension collapses to the GQA factor (16), so the MMA tile is 16×128 — decode-shaped, and naturally served by, e.g., a FlashInfer decode attention kernel. Tensor cores are underfed, and a KV block selected by M queries is read M times repeatedly. • KV-outer (KV-stationary). Invert the map: for each (kv_head, kv_block), gather the queries that selected it, load the KV block once, and run a full 128×128 tile, emitting one partial (O_partial, m, l) per (query, kv_block). So each Q token writes up to topK (= 16) partial-O blocks, which a following combine kernel LSE-merges into its final output. Full tiles; each KV block read once.

Figure: kernel scheduling comparison between FA4 (dense), Q-outer (sparse) and KV-outer (sparse). Left: FA4 (dense); compute: efficient, 128x128 MMA; load: efficient, each KV block is shared by 128 Q tokens; store: efficient, O is stored once. Middle: Q-outer (sparse); compute: inefficient, 16 (MQA factor) x128 MMA; load: inefficient, each KV block is loaded by multiple Q tokens separately; store: efficient, O is stored once. Right: KV-outer (sparse); compute: efficient, 128x128 MMA; load: okay, each Q block is shared by 128 KV tokens via gathered read; store: inefficient, partial-O is stored 16 (topk) times. I/O analysis: sparse Q-outer vs sparse KV-outer crossover

While Section 4.2 of the Minimax Sparse Attention paper analyzes the FLOPs/IO tradeoffs between Q-outer and KV-outer approaches, it overlooks the fact that most data movement occurs within the L2 cache rather than global memory, necessitating a distinct analytical approach. The two structures are bound by different memories. Q-outer is L2-bandwidth bound (re-reads hit L2); KV-outer is HBM-bandwidth bound (each block read once, but ≤topK partials are written per query and re-read in combine). A simple bandwidth roofline (B200 at ~24 TB/s L2 and ~7.4 TB/s HBM, from measured peak numbers; data fp8, partial-O bf16) gives the crossover: KV-outer is faster than Q-outer when nsb / N HBM latency ⇒ L2 - bw bound

KV - outer HBM ( dedup ) : N· ( Q + O + partial - O ) + nsb· ( K + V ) = 69 , 632 ·N + 32 , 768 ·nsb latency = ( 69 , 632 ·N + 32 , 768 ·nsb ) / 7.4 = 9 , 410 ·N + 4 , 428 ·nsb L2 ( re - read ) : N·topK·Q + 2 ·N·partial - O + N·O + nsb· ( K + V ) = 165 , 888 ·N + 32 , 768 ·nsb latency = ( 165 , 888 ·N + 32 , 768 ·nsb ) / 24 = 6 , 912 ·N + 1 , 365 ·nsb HBM latency > L2 latency ( always ) ⇒ HBM - bw bound

roofline comparison ( each at its binding resource ) Q - outer ( L2 ) = 22 , 016 ·N KV - outer ( HBM ) = 9 , 410 ·N + 4 , 428 ·nsb KV - outer faster : 22 , 016 ·N > 9 , 410 ·N + 4 , 428 ·nsb ⇒ nsb / N < 2.85

Caveats. (1) The bound is conservative — it gives Q-outer a perfect L2 hit; beyond L2 capacity Q-outer is HBM-bound and slower, so KV-outer wins earlier than 2.85. (2) Bandwidth-only — tensor-core and combine compute are not modeled,...