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

def mla_prefill[q_type: DType, k_t: MHAOperand, v_t: MHAOperand, k_rope_t: MHAOperand, output_type: DType, mask_t: MHAMask, valid_layout: TensorLayout, config: MHAConfig[config.dtype], group: Int = Int(128), q_depth: Int = Int(192), cache_depth: Int = Int(576), _ndbuffer_mha_operand: Bool = False](q_ptr: Pointer[Scalar[q_type], MutAnyOrigin, _safe=False], k: k_t, v: v_t, k_rope: k_rope_t, output_ptr: Pointer[Scalar[output_type], MutAnyOrigin, _safe=False], scale: Float32, batch_size: Int, seq_len_arg: Int, valid_length_tt: TileTensor[DType.uint32, valid_layout, ImmutAnyOrigin], cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]], mask: mask_t)

MLA prefill kernel that computes attention for sequences longer than one token.

Resolves per-batch sequence lengths from the ragged row-offsets, then delegates to mla_prefill_single_batch on NVIDIA or Attention.mla_prefill on AMD to compute the QK and PV matmuls with online softmax over the full prompt plus cached KV.

Parameters:

q_type (DType): Element type of the query tensor q_ptr (inferred). Must be a half-float or float8_e4m3fn.
k_t (MHAOperand): Key operand type backing k (inferred). Provides the nope part of K.
v_t (MHAOperand): Value operand type backing v (inferred).
k_rope_t (MHAOperand): KV cache operand type backing k_rope (inferred). Provides the rope part of K and the per-batch cache length.
output_type (DType): Element type of output_ptr (inferred). bfloat16 when q_type is float8_e4m3fn, else matches q_type.
mask_t (MHAMask): Mask functor type applied to attention scores (inferred).
valid_layout (TensorLayout): Compile-time layout of valid_length_tt (inferred).
config (MHAConfig[config.dtype]): MLA attention config carrying num_heads, depth, and tile geometry (block_m, block_n, block_k).
group (Int): GQA group size, num_heads // num_kv_heads (defaults to 128).
q_depth (Int): Q head dimension in elements (defaults to 192). For DeepSeek V2/3, 192 = 128 nope + 64 rope.
cache_depth (Int): KV cache head dimension in elements (defaults to 576). The absorbed-latent width of the paged MLA cache.
_ndbuffer_mha_operand (Bool): Whether the K/V/k_rope operands use ND-buffer layout instead of the default TileTensor layout (defaults to False).

Args:

q_ptr (Pointer[Scalar[q_type], MutAnyOrigin, _safe=False]): Pointer to the query tensor with shape [total_q_tokens, num_heads, q_depth].
k (k_t): Key operand (nope part).
v (v_t): Value operand.
k_rope (k_rope_t): KV cache operand providing the rope part of K; k_rope.cache_length(batch_idx) gives the cached length.
output_ptr (Pointer[Scalar[output_type], MutAnyOrigin, _safe=False]): Pointer to the output tensor with shape [total_q_tokens, num_heads, depth].
scale (Float32): Softmax scale factor applied to QK^T.
batch_size (Int): Number of batches in the request; drives the grid's z dimension.
seq_len_arg (Int): Maximum query sequence length (padding ceiling across all batches).
valid_length_tt (TileTensor[DType.uint32, valid_layout, ImmutAnyOrigin]): Per-batch uint32 tensor of cumulative row offsets with batch_size + 1 entries; offsets[b] to offsets[b+1] gives batch b's token range.
cache_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]]): Optional per-batch uint32 tensor of starting offsets into the paged KV cache; None when the cache is contiguous.
mask (mask_t): Mask functor instance applied to attention scores.