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_prefill_branch_sparse_fp8

def mla_prefill_branch_sparse_fp8[dtype: DType, fp8_dtype: DType, fp8_scale_dtype: DType, collection_t: KVCollectionT, //, m_scale_granularity: Int, n_scale_granularity: Int, k_scale_granularity: Int, kv_input_fn: def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[DType.bfloat16, width], indices_stride: Int, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu")](output: TileTensor[dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type], freqs_cis: TileTensor[Storage=freqs_cis.Storage, linear_idx_type=freqs_cis.linear_idx_type], kv_norm_gamma: TileTensor[Storage=kv_norm_gamma.Storage, linear_idx_type=kv_norm_gamma.linear_idx_type], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, epsilon: Float32, w_uk: TileTensor[fp8_dtype, Storage=w_uk.Storage, linear_idx_type=w_uk.linear_idx_type], w_uk_scale: TileTensor[fp8_scale_dtype, Storage=w_uk_scale.Storage, linear_idx_type=w_uk_scale.linear_idx_type], w_uv: TileTensor[fp8_dtype, Storage=w_uv.Storage, linear_idx_type=w_uv.linear_idx_type], w_uv_scale: TileTensor[fp8_scale_dtype, Storage=w_uv_scale.Storage, linear_idx_type=w_uv_scale.linear_idx_type], ctx: DeviceContext, d_indices: Pointer[Int32, MutAnyOrigin, _safe=False], topk_lengths: Pointer[Int32, MutAnyOrigin, _safe=False], attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]])

Sparse MLA prefill branch (DSv3.2/GLM absorbed shape, FP8 weights).

Reuses mla_decode_branch_fp8's absorbed-Q construction (q_nope up-proj via w_uk + RoPE/RMSNorm) and the identical w_uv output up-projection, and swaps the attention call to the existing mla_sm100_prefill_sparse kernel over the paged BF16 latent cache. The caller (.fp8.sparse op) has already remapped d_indices from logical to physical rows, so they are passed straight through. Only supported for a BF16 KV cache; the FP8-cache case is handled by the caller (dense fallback).