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_blockscale

def mla_sm100_prefill_blockscale[output_dtype: DType, q_type: DType, KVType: MHAOperand, VType: MHAOperand, KRopeType: MHAOperand, MaskType: MHAMask, MaxPromptLenType: OptionallyStaticInt, //, config: MHAConfig[config.dtype], group: Int, q_depth: Int, cache_depth: Int, _ndbuffer_mha_operand: Bool, blockwise_scale: Int = Int(0), v_depth: Int = Int(-1)](output: TileTensor[output_dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[q_type, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: KVType, v: VType, k_rope: KRopeType, mask_functor: MaskType, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], max_prompt_len: MaxPromptLenType, scale: Float32, batch_size: Int, ctx: DeviceContext)

Launches the SM100 MLA prefill kernel with blockwise FP8 scaling.

Selects an FA4 MLA config, builds the TMA tiles for Q, K_nope, K_rope, and V, then enqueues the warp-specialized mla_prefill_kernel_blockscale kernel that converts block-scaled FP8 K_rope to BF16 before the MMA stage.

Parameters:

output_dtype (DType): The element dtype of the output tensor.
q_type (DType): The element dtype of the query tensor.
KVType (MHAOperand): The MHAOperand type for K_nope, carrying gmem layout and page-size metadata.
VType (MHAOperand): The MHAOperand type for V, carrying gmem layout and page-size metadata.
KRopeType (MHAOperand): The MHAOperand type for K_rope, carrying gmem layout, page-size, and dtype (FP8 for blockscale).
MaskType (MHAMask): The MHAMask type controlling causal and padding behavior.
MaxPromptLenType (OptionallyStaticInt): The OptionallyStaticInt type of max_prompt_len, either a compile-time constant or a runtime integer.
config (MHAConfig[config.dtype]): The MHAConfig holding tile sizes, stage counts, and head counts (inferred).
group (Int): The GQA group size, namely the number of query heads per KV head (inferred).
q_depth (Int): The total query head dimension, combining nope and rope depth (inferred).
cache_depth (Int): The total KV cache depth, combining nope and rope depth (inferred).
_ndbuffer_mha_operand (Bool): Whether the MHA operand uses the ndbuffer layout (inferred).
blockwise_scale (Int): The blockwise FP8 scaling mode flag (defaults to 0, inferred).
v_depth (Int): The V and output head dimension, or -1 for the DeepSeek shape where V width equals nope width (defaults to -1, inferred).

Args:

output (TileTensor[output_dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): The output tensor of shape [batch * num_rows, num_q_heads, ov_depth].
q (TileTensor[q_type, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): The query tensor of shape [batch * num_rows, num_q_heads, qk_depth].
k (KVType): The K_nope operand for the non-rotary part of the KV cache.
v (VType): The V operand for the value part of the KV cache.
k_rope (KRopeType): The K_rope operand for the rotary-position-encoded part of K, stored as block-scaled FP8.
mask_functor (MaskType): The attention mask functor applied to the score matrix.
valid_length (TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type]): The per-sequence valid length tensor of shape [batch] as uint32.
max_prompt_len (MaxPromptLenType): The maximum prompt length across the batch, as a compile-time or runtime integer.
scale (Float32): The softmax scale factor applied to the Q@K' dot product.
batch_size (Int): The number of sequences in the batch.
ctx (DeviceContext): The DeviceContext used to enqueue the kernel.