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

def mla_sm100_prefill[output_type: 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_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], 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)

Dense MLA prefill dispatcher for SM100 (Blackwell).

Routes to mla_sm100_prefill_generic when blockwise_scale is zero and the query, key, value, and rope dtypes all match, otherwise routes to mla_sm100_prefill_blockscale. Asserts that the output is BF16 and that the key and value share an element dtype.

Parameters:

output_type (DType): Output element type (must be DType.bfloat16).
q_type (DType): Query element type.
KVType (MHAOperand): Key operand type (an MHAOperand).
VType (MHAOperand): Value operand type (an MHAOperand).
KRopeType (MHAOperand): Rope key operand type (an MHAOperand).
MaskType (MHAMask): Attention mask functor type (an MHAMask).
MaxPromptLenType (OptionallyStaticInt): Maximum prompt length type (an OptionallyStaticInt).
config (MHAConfig[config.dtype]): MHA configuration struct.
group (Int): Number of query heads per KV head (GQA group size).
q_depth (Int): Per-head query depth.
cache_depth (Int): Per-head KV cache depth.
_ndbuffer_mha_operand (Bool): Whether operands are ND buffers.
blockwise_scale (Int): Blockwise quantization scale size; zero disables blockwise scaling and selects the generic path.
v_depth (Int): Per-head value depth; -1 resolves to the nope width (DeepSeek convention) inside MLAConfig.__init__.

Args:

output (TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Output tile tensor with shape [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.
k (KVType): Key operand.
v (VType): Value operand.
k_rope (KRopeType): Rope-applied key operand.
mask_functor (MaskType): Attention mask functor.
valid_length (TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type]): Per-sequence valid lengths (uint32).
max_prompt_len (MaxPromptLenType): Maximum prompt length (static or dynamic).
scale (Float32): Softmax scale.
batch_size (Int): Number of sequences in the batch.
ctx (DeviceContext): GPU device context.