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
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.
Each lane owns a set of (key_row, (token, head)_col) fragment elements
(identical ownership to the MSA epilogue at BM = 64). The epilogue applies the
branchless relu (x + |x|) * 0.5, multiplies by q_scale[token, head], sums
over the head columns of each token (columns of the same token in the N-tile),
cross-lane reduces over the 4 lanes sharing a key row (lane_group_sum, changed
from the MSA's max), multiplies by k_scale[key], and writes one f32 per
(token, key).
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 = 64 (256 CTAs at batch=8, num_keys=2048) intentionally oversubscribes
the grid for the launch-starved decode regime and reuses the MSA BM = 64
fragment map verbatim. 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β
Functionsβ
- β
fp8_index_score_sm100: Launch the SM100 tensor-core FP8 indexer scorer intooutput.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!