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_single_batch

def mla_prefill_single_batch[q_type: DType, k_t: MHAOperand, v_t: MHAOperand, k_rope_t: MHAOperand, output_type: DType, mask_t: MHAMask, *, config: MHAConfig[config.dtype], group: Int = Int(1), q_depth: Int = Int(192), cache_depth: Int = Int(576)](q_ptr: Pointer[Scalar[q_type], _safe=False], k: k_t, v: v_t, k_rope: k_rope_t, output_ptr: Pointer[Scalar[output_type], _safe=False], scale: Float32, seq_len: Int, max_seq_len: Int, start_pos: UInt32, cache_start_pos: UInt32, num_keys: Int, mask: mask_t, batch_idx: Int)

MLA for encoding where seqlen > 1.

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).
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 1).
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.

Args:

q_ptr (Pointer[Scalar[q_type], _safe=False]): Pointer to this batch's query rows with shape [seq_len, num_heads, q_depth].
k (k_t): Key operand providing the nope (latent) part of K.
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], _safe=False]): Pointer to this batch's output rows with shape [seq_len, num_heads, depth] where depth is config.depth.
scale (Float32): Softmax scale factor applied to QK^T.
seq_len (Int): Valid sequence length (without padding) for this batch.
max_seq_len (Int): Padded sequence length ceiling for this batch.
start_pos (UInt32): Starting position of the query within the sequence (the cached length before this prefill).
cache_start_pos (UInt32): Starting offset into the paged KV cache for this batch.
num_keys (Int): Total number of KV keys (cached plus new tokens) for this batch.
mask (mask_t): Mask functor instance applied to attention scores.
batch_idx (Int): Index of the current batch in the request.