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

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 BM_key-key tile and streams every query token past it. That maps a batch-1 prefill onto only num_keys / BM_key CTAs 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:

• 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 = 128, 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), num_key_parts): one CTA per (query-token block, key part). A batch-1 GLM prefill (1024 tokens, num_heads=32, N_TOKENS=4) is 256 CTAs on the first two axes alone, so it runs unsplit at num_key_parts == 1. Decode/MTP inverts that -- a handful of token blocks over a long cache -- and grid.z supplies the parallelism instead, each CTA streaming _KEY_TILES_PER_CTA tiles of its own key window. num_key_parts is a grid EXTENT, not the realized split: it is sized from the batch maximum, so on a ragged batch each CTA narrows it to what its own entry can feed (_MIN_TILES_PER_PART) and the surplus parts retire immediately.

Warp specialization, mirroring the MSA prefill scorer (Kernels/lib/msa/sparse_indexer_prefill.mojo, PR #91938), on a 256-thread CTA at every tile:

• WG0 (warps 0-3, threads 0-127) = score/epilogue consumer. Drains each S^T stage out of TMEM one row per thread (tcgen05.ld.32x32b, warp w / lane l -> row 32 * w + l, so the 4 warps span all BM_key rows), applies the branchless relu, sums over each token's head columns entirely within the thread, scales by k_scale, and writes one f32 per (token, key) under the fused causal guard.
• WG1 (warps 4-7) = producer: warp 4 = MMA (TMEM owner + K @ Q^T per K tile), warp 5 = TMA (deep K-ring producer). Warps 6-7 are idle and register-dealloc to the floor; they exist only so setmaxnreg has a whole warpgroup to issue the dec from, which is what funds the consumer's 216. Role-to-role mbars (k_full/k_empty for the K ring, s_full/s_empty for the multi-stage S^T) replace the shipped kernel's per-iteration whole-CTA named_barrier; the resident Q owns no barrier and rides k_full[0].

Scores carry no cross-key reduction (no softmax denominator), and a thread's only global write is output[global_token, key_local]. Key windows are therefore disjoint output elements: the split needs no combine pass, no workspace, and no atomics.

Routing lives in fp8_index_score_sm100; see the comment there for the measured thresholds. Two disjoint corners land here: enough token blocks to fill the grid on their own (long prefill), or too few token blocks but a key range deep enough that splitting it fills the grid (decode / MTP-decode). The kernel body supports all head counts uniformly; the route currently admits num_heads in {32, 64}.

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.