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
SM100 (B200) tensor-core FP8 MLA lightning-indexer score kernel.
Computes the same per-(query token, key) logit as the scalar
nn.index_fp8.fp8_index_kernel, but runs the depth-128 dot product on the
tcgen05 tensor cores instead of a serial FMA loop:
score[token, key] = k_scale[key]
* Σ_head relu(q[token, head] · k[key]) * q_scale[token, head]Q is [total_seq, num_heads, depth] fp8-e4m3, K is paged [keys, depth] fp8-e4m3
with a per-token k_scale, and the head reduction is a sum over num_heads.
Layout (crux), cloning the shipped MSA prefill scorer
(Kernels/lib/msa/sparse_indexer_prefill.mojo) with the operand roles inverted:
- MMA_M = key tile (
BM_key): A operand = this CTA's K tile[BM_key, depth], loaded once and reused across the MTP query tokens. - MMA_N = 128 = (query-token × head): B operand = a pair of query tokens'
[N_TOKENS * num_heads, depth],transpose_b=True->S^T = K @ Q^T = [key, (token, head)]. This is DeepGEMM'ssm100_mqa_logitspacking (BLOCK_Q = 128 / num_heads, soN_TOKENS = 2atnum_heads = 64). - MMA_K = depth = 128 contraction, fp8 in / f32 TMEM accumulation.
The epilogue drains TMEM one row per thread (tcgen05.ld.32x32b, warp w /
lane l -> accumulator row 32 * w + l), so a thread owns a whole key row and
the head reduction never leaves its registers. Per column it applies the
branchless relu (x + |x|) * 0.5 and multiplies by q_scale[token, head],
summing over each token's head columns; then it multiplies by
k_scale[key] and writes one f32 per (token, key). Because a warp holds 32
consecutive keys at a fixed token, each store is one fully-active 128B
transaction.
Grid (batch, ceil(num_keys / BM_key), seq_slices): key tiles and token tiles
are independent outputs, so there is no split-K and no cross-CTA reduction.
BM_key = MMA_M = 128 keeps the standard (non-.ws) tcgen05 datapath -- the
packed .ws form engages at MMA_M <= 64, and the non-ws form there would
leave half the datapaths idle -- and matches the 128 TMEM lanes this CTA's four
warps drain. grid.z splits a sequence's token tiles across CTAs when the
key-tile grid alone underfills the machine (low-key prefill); decode always
launches one slice.
Token tiles are software-pipelined: Q and its scales are double-buffered so tile nt+1's TMA and q_scale loads fly under tile nt's MMA and epilogue. The TMEM accumulator stays single-stage (the drain is a TMEM->register copy that precedes the epilogue math, so the next MMA already overlaps the math; a second stage measured as a pure loss by halving TMEM-limited CTAs/SM on large grids). Decode launches allocate only the SMEM prefix (Q buffer 1 is last in the layout and unreachable at a single token tile), keeping decode occupancy unchanged.
Prefill / causal masking and the -inf tail-fill fusion are Slice 2/3 (not here);
this kernel is a drop-in for the score buffer the top-k stage consumes.
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
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
SPEC_DECODE_N_TOKENS_ALT
comptime SPEC_DECODE_N_TOKENS_ALT = 3
Functions
-
fp8_index_score_sm100: Launch the SM100 tensor-core FP8 indexer scorer intooutput.