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_prefill_sparse_qkv_fp8
SM100 (B200) native-FP8 sparse MLA prefill kernel.
Two comptime-selected tiles share this file: head<=64 on a single CTA
(cta_group=1) with K and V sharing one k-major gather buffer, and head=128 on a
2-CTA cluster (cta_group=2) with CTA-split K + a separate full-V gather (the
shared-KV simplification is impossible when CTA-split-K@base0 and full-V@base0
collide). The cg2 cluster scaffold (grid, cta_id/seq_idx split, tcgen05_alloc
[cta_group], CTA_MASK multicast, cross-CTA arrive_cluster handshakes) mirrors the
dequant mla_prefill_sparse_kv_fp8.mojo; the compute stays native f8f6f4. Both
tiles read V mn-major from a k-major SW64 buffer at BK=b_topk=64.
All of Q, K, V, and P are FP8 e4m3; QK^T and P*V run natively on tensor cores
via tcgen05.mma.kind::f8f6f4 (KIND_F8F6F4) with NO in-SMEM FP8->BF16 dequant.
This removes the convert_k/v_fp8_to_bf16 stage (and its LOCAL register staging)
of mla_prefill_sparse_kv_fp8.mojo, so the kernel spills zero registers by
construction and holds K/V at 1 byte/element.
Structure (top-k gather, 4-warpgroup layout, launch) mirrors
mla_prefill_sparse_kv_fp8.mojo. Native-FP8 MMA machinery (single SS QK^T,
2-atom P*V with V read mn-major from the same k-major KV buffer, FP8 P written to
a separate SMEM region) mirrors the decode reference mla_decode_qkv_fp8.mojo
(DecodeSM100QKTSS_FP8 / DecodeSM100PVSS_FP8). Both operand-layout
foundations were validated bit-exact on B200; see
.claude/agent-memory/mojo-kernel-engineer/fp8-sw64-gather4-qk-mma-validated.md.
SMEM layout (all FP8 operands, num_mbars=2 stages): Q FP8: 64 x 576 x 1 = 36864 bytes (persistent, SWIZZLE_64B) KV FP8: 2 x 64 x 576 x 1 = 73728 bytes (double-buffered, SWIZZLE_64B) P FP8: 2 x 64 x 64 x 1 = 8192 bytes (double-buffered, SWIZZLE_64B) O bf16: reuses the KV region at the epilogue
The KV buffer is the single k-major SW64 gather; V is read mn-major from its
first v_depth columns (K and V share the MLA latent, so there is no separate V
gather). Scale handling targets the DSv3.2 production scale_block_size == 0
(unit-scale) path: kv_scale == v_scale == 1, so scale_log2e and output_scale
match the BF16-KV kernel's unit-scale formulas. Tensorwise/blockwise scale-fold
and head=128 (cta_group=2, 2-CTA f8f6f4) are follow-ups.
comptime values
FP8_TYPE
comptime FP8_TYPE = DType.float8_e4m3fn
SW64
comptime SW64 = TensorMapSwizzle.SWIZZLE_64B
Structs
-
MLAPrefillSparseQKVFP8: -
MLASparseSharedMemoryQKVFP8: Native-FP8 SMEM layout: FP8 Q/KV/P + FP8-agnostic softmax scratch.
Functions
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!