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

sparse_index_fp8_sm100_prefill

SM100 (B200) warp-specialized PREFILL variant of the FP8 MLA indexer scorer.

Computes the identical per-(query token, key) logit as the shipped scorer (sparse_index_fp8_sm100.fp8_index_score_sm100) and the scalar nn.index_fp8.fp8_index_kernel:

score[token, key] = k_scale[key]
                    * Σ_head relu(q[token, head] · k[key]) * q_scale[token, head]

The shipped kernel is K-resident / Q-streaming: one CTA holds a 64-key tile and streams every query token past it. That maps a batch-1 prefill onto only num_keys / 64 CTAs (128 at 8192 keys < 148 SMs) and runs one serial warpgroup, so it is latency-bound (measured ~6% achieved occupancy: one active warp per scheduler, no MMA↔epilogue overlap).

This kernel INVERTS which operand persists (the S^T orientation and the BM=64 head-sum fragment map are kept verbatim -- only the streamed operand changes):

• Q resident as the B operand [MMA_N = N_TOKENS * num_heads, depth]: a CTA owns one N_TOKENS-token block, staged once.
• K streams as the A operand [BM_key = 64, depth] through a deep SMEM prefetch ring, S^T = K @ Q^T = [key, (token, head)], so the epilogue reduces over the (token, head) COLUMNS exactly like the shipped kernel (heads stay columns; all head counts in {4, 8, 32, 64} work uniformly, no cross-warp reduction).
• Grid (batch, ceil(seq_len / N_TOKENS), 1): one CTA per query-token block. A batch-1 GLM prefill (1024 tokens, num_heads=32, N_TOKENS=4) is 256 CTAs -- it fills the machine without a seq-split -- and each CTA streams every key.

Warp specialization (256 threads = two warpgroups), mirroring the MSA prefill scorer (Kernels/lib/msa/sparse_indexer_prefill.mojo, PR #91938):

• WG0 (warps 0-3, threads 0-127) = score/epilogue consumer. Reads each S^T stage out of TMEM, applies the branchless relu, sums over the head columns of each token, cross-lane reduces the 4 lanes sharing a key row, scales by k_scale, and writes one f32 per (token, key) under the fused causal guard.
• WG1 (warps 4-7, threads 128-255) = producer: warp 4 = MMA (TMEM owner + K @ Q^T per K tile), warp 5 = TMA (deep K-ring producer), warps 6-7 idle (register-dealloc to the floor). Role-to-role mbars (k_full/k_empty for the K ring, the accumulator's own MMA↔TMEM handshake for the multi-stage S^T) replace the shipped kernel's per-iteration whole-CTA named_barrier.

Token blocks are independent outputs, so there is no split-K and no cross-CTA reduction. Because a CTA sees EVERY key for its resident token block, this layout is also the prerequisite for the planned fused score+top-k-select variant that never materializes the [total_seq, max_num_keys] score buffer (see KERN-3139).

Standalone routing (measured): the token-block grid is ceil(seq / N_TOKENS) with N_TOKENS = 128 // num_heads, so it only fills the machine -- and only wins over the K-resident scorer -- at num_heads == 64 (N_TOKENS = 2; measured +7-17% at nh64 prefill). At num_heads in {4, 8, 32} too few token blocks underfill the grid and the heavier per-key-tile head-sum makes the kernel consumer-bound (a net loss vs K-resident), so fp8_index_score_sm100 routes only num_heads == 64 prefill here; every other shape (and decode) stays on the shipped K-resident path byte-identically. The kernel body itself supports all head counts uniformly, kept general for the fused score+top-k-select follow-up.

NVIDIA SM100 only (SS-UMMA / TMA / tcgen05). Verified against nn.index_fp8.fp8_index_naive via test_index_fp8 and end-to-end top-k set match via test_mla_index_fp8.

comptime values

KTMATileT

comptime KTMATileT[dtype: DType, BM_key: Int, depth: Int] = TMATensorTile[dtype, Int(3), _padded_shape[Int(3), dtype, Index[Int, Int, Int](BM_key, Int(1), depth), TensorMapSwizzle.SWIZZLE_128B](), _ragged_shape[Int(3), dtype, Index[Int, Int, Int](BM_key, Int(1), depth), TensorMapSwizzle.SWIZZLE_128B]()]

Parameters

• ​dtype (DType):
• ​BM_key (Int):
• ​depth (Int):

QTMATileT

comptime QTMATileT[dtype: DType, MMA_N: Int, depth: Int] = TMATensorTile[dtype, Int(3), _padded_shape[Int(3), dtype, Index[Int, Int, Int](MMA_N, Int(1), depth), TensorMapSwizzle.SWIZZLE_128B](), _ragged_shape[Int(3), dtype, Index[Int, Int, Int](MMA_N, Int(1), depth), TensorMapSwizzle.SWIZZLE_128B]()]

Parameters

• ​dtype (DType):
• ​MMA_N (Int):
• ​depth (Int):

Functions

• ​fp8_index_score_sm100_prefill: Enqueue the warp-specialized K-streaming prefill scorer into output.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit