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

mla_decode_sparse_qkv_fp8

Native FP8 sparse MLA decode kernel for SM100 (B200) -- unified gather.

Combines the OLD kernel's efficient contiguous KV gather (mla_decode_sparse_kv_fp8.mojo: INT64/SWIZZLE_NONE gather4, 16 gather4 instructions/tile at 2304 B/txn) with the native-FP8 compute core (tcgen05.mma.kind::f8f6f4, no dtype conversion). The two are decoupled by a re-swizzle staging warpgroup: KV gathers contiguously into a LINEAR FP8 SMEM buffer, then WG3 (repurposing the old kernel's convert-WG role, minus the dtype cast) permutes it FP8->FP8 into the SW64 byte layout the FP8 tcgen05 operand descriptor expects, into a SECOND SMEM buffer. QK^T and P*V still run natively in FP8 reading the SW64 buffer -- only the KV front-end changed; the compute core is byte-for-byte what MLA_SM100_Decode_Sparse_QKV_FP8 already validated.

Root cause this fixes (see .claude/agent-memory/mojo-kernel-engineer/ mla-decode-sm100-native-fp8-warpgroup-count.md): the native kernel's own K/Q operand layout mandates SWIZZLE_64B for tcgen05.mma.kind::f8f6f4, and a gather4 TMA under SW64 has box_width=64B, so a 576-byte row needs ceildiv(576,64)=9 column-groups x 16 four-row-chunks = 144 gather4 instructions/tile -- 9x the old kernel's 16 instructions/tile (SWIZZLE_NONE, single 576-byte box) for identical bytes. That per-tile hardware-instruction inflation, not pipeline depth or warp count, is what NCU's PC-sampled stall attribution pinned to the KV-gather producer path.

4 warpgroups (back from 3 -- the re-swizzle stage, like the old kernel's convert stage, needs its own resident warps; it is a pure FP8->FP8 layout permute, not a dtype conversion): WG0 (warps 0-3): Softmax (native_fp8=True) WG1 (warps 4-7): Correction WG2 (warps 8-11): Load (warp 8, INT64/SWIZZLE_NONE gather4) / MMA QK (warp 9) / MMA PV (warp 10) / idx_producer + Store (warp 11) -- unchanged roles from the 3-WG kernel. WG3 (warps 12-15): Re-swizzle: read the contiguously-gathered FP8 linearly, write FP8 (unchanged bytes, no cast) into the SW64 layout via a manual swizzled SMEM store (all 128 threads cooperate, mirroring convertFP8ToBF16's thread mapping but without the dtype cast or the BF16-overlay aliasing -- this uses a separate, non-overlaid destination buffer).

SMEM layout (FP8, N = num_kv_stages; expect N ~= 2 at the GLM-5.2 shape, since the linear + SW64 buffer pair costs the same total bytes/stage as the old kernel's own KV region): Q FP8: 64 x 576 x 1 = 36864 bytes (SWIZZLE_64B, persistent) KV linear-in: N x 64 x 576 x 1 bytes (SWIZZLE_NONE, gather dest) KV SW64-out: N x 64 x 576 x 1 bytes (SWIZZLE_64B, MMA operand) P stages: N x 64 x 64 x 1 bytes (SWIZZLE_64B, separate region) max/li: 128 x 4 x 3 = 1536 bytes barriers: (8N+11) fixed + output + idx_bars (2N) barriers ptr_tmem_addr: 4 bytes idx_smem: N x 64 x 4 bytes

Structs

• ​MLA_SM100_Decode_Sparse_QKV_FP8: Sparse MLA decode with native FP8 WGMMA for SM100 (B200), unified gather.