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 function

mla_sm100_prefill_sparse_fp8

def mla_sm100_prefill_sparse_fp8[output_type: DType, q_type: DType, cache_t: KVCacheT, //, num_q_heads: Int, qk_depth: Int, v_depth: Int, indices_stride: Int, scale_block_size: Int](output: TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[q_type, Storage=q.Storage, linear_idx_type=q.linear_idx_type], kv_cache: cache_t, indices: TileTensor[DType.uint32, Storage=indices.Storage, linear_idx_type=indices.linear_idx_type], topk_lengths: TileTensor[DType.uint32, Storage=topk_lengths.Storage, linear_idx_type=topk_lengths.linear_idx_type], attn_sink_ptr: Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]], scales_ptr: Pointer[Float32, ImmutAnyOrigin, _safe=False], scale: Float32, ctx: DeviceContext)

FP8 KV-cache variant of mla_sm100_prefill_sparse.

Thin wrapper around mla_prefill_sparse_fp8 that builds the MLASparseConfig from the passed dimensions. The FP8 KV cache uses DType.float8_e4m3fn storage with one Float32 scale per physical KV token supplied via scales_ptr.

Parameters:

output_type (DType): Output element type (BF16).
q_type (DType): Query element type (BF16).
cache_t (KVCacheT): FP8 KV cache type (DType.float8_e4m3fn).
num_q_heads (Int): Number of query heads (128 for DSv3.2).
qk_depth (Int): Per-head Q/K depth (576 for DSv3.2).
v_depth (Int): Per-head V depth (512 for DSv3.2).
indices_stride (Int): Per-query indices buffer stride (= top-k count).
scale_block_size (Int): Quantization block size along the depth axis. 0 selects no-scale mode: FP8 latents are read at unit scale (no dequant, scales_ptr unused / may be null), mirroring the sparse-decode read of today's scale-less MLA latent cache — the current DSv3.2 production path. >= qk_depth selects tensorwise scaling (one scale per KV token). A smaller positive value (e.g. 32) selects cache-native blockwise scaling: MLA stores one ceildiv(qk_depth, scale_block_size)-wide scale vector per token over the full latent, K consumes all of it and V the first ceildiv(v_depth, scale_block_size) blocks, so a single scales_ptr (stride ceildiv(qk_depth, sbs)) serves both.

Args:

output (TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Output tile tensor [total_q_tokens, num_q_heads, v_depth].
q (TileTensor[q_type, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query tile tensor [total_q_tokens, num_q_heads, qk_depth].
kv_cache (cache_t): Paged MLA FP8 KV cache for the current layer.
indices (TileTensor[DType.uint32, Storage=indices.Storage, linear_idx_type=indices.linear_idx_type]): Per-query gather4 indices (uint32).
topk_lengths (TileTensor[DType.uint32, Storage=topk_lengths.Storage, linear_idx_type=topk_lengths.linear_idx_type]): Per-query effective top-k count.
attn_sink_ptr (Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]]): Optional attention sink (pass None to skip).
scales_ptr (Pointer[Float32, ImmutAnyOrigin, _safe=False]): FP8 dequantization scales in get_tma_row space, laid out [total_phys_rows, ceildiv(qk_depth, scale_block_size)] (i.e. k.scales_raw_ptr() for the paged cache).
scale (Float32): Softmax scale.
ctx (DeviceContext): GPU device context.