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
mla_decode_sparse_qkv_fp8
Native FP8 sparse MLA decode kernel for SM100 (B200).
Combines the sparse gather4 TMA machinery from mla_decode_sparse_kv_fp8.mojo with the native FP8 3-WG pipeline from mla_decode_qkv_fp8.mojo.
Eliminates the FP8→BF16 convert warpgroup (WG3) present in MLA_SM100_Decode_Sparse_KV_FP8 by performing QK and PV MMAs natively in FP8 (tcgen05.mma.kind::f8f6f4). Q arrives as FP8 via SWIZZLE_64B TMA directly. KV arrives as FP8 via gather4 INT64 TMA directly into kv_smem (no SMEM conversion stage). P is FP8 in a separate SMEM region.
Going from 4 warpgroups back to 3: WG0 (warps 0-3): Softmax (native_fp8=True) WG1 (warps 4-7): Correction WG2 (warps 8-11): Load (warp 8) / MMA QK (warp 9) / MMA PV (warp 10) / idx_producer + Store (warp 11)
SMEM layout (FP8, N = num_kv_stages, typically 4): Q FP8: 64 x 576 x 1 = 36864 bytes (SWIZZLE_64B) KV stages: N x 64 x 576 x 1 bytes (SWIZZLE_64B) P stages: N x 64 x 64 x 1 bytes (SWIZZLE_64B, separate region) max/li: 128 x 4 x 3 = 1536 bytes barriers: (6N+11) fixed + output + idx_bars (4) barriers ptr_tmem_addr: 4 bytes idx_smem: 2 x 64 x 4 = 512 bytes
Structs
-
MLA_SM100_Decode_Sparse_QKV_FP8: Sparse MLA decode with native FP8 WGMMA for SM100 (B200).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!