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_decoding

def mla_decoding[q_type: DType, k_t: MHAOperand, output_type: DType, mask_t: MHAMask, ValidLT: TensorLayout, BM: Int, BN: Int, BK: Int, WM: Int, WN: Int, depth: Int, depth_v: Int, num_heads: Int, num_threads: Int, num_pipeline_stages: Int, group: Int = Int(1), ragged: Bool = False, _use_valid_length: Bool = False, _is_cache_length_accurate: Bool = False, decoding_warp_split_k: Bool = False, q_seq_len: Int = Int(1)](q_ptr: Pointer[Scalar[q_type], MutAnyOrigin, _safe=False], k: k_t, output_ptr: Pointer[Scalar[output_type], MutAnyOrigin, _safe=False], exp_sum_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False], qk_max_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False], scale: Float32, batch_size: Int, num_partitions: Int, max_cache_valid_length: Int, valid_length_tt: TileTensor[DType.uint32, ValidLT, ImmutAnyOrigin], mask: mask_t)

MLA decoding kernel that computes single-batch attention per CTA.

Sets up split-K offsets, batch indexing, and ragged or dense query strides, then delegates to mla_decoding_single_batch on NVIDIA or Attention.mla_decode on AMD to perform the actual QK and PV computation for one decode step (with optional MTP query folding on AMD).

Parameters:

q_type (DType): Element type of the query tensor q_ptr. Must be a half-float or float8_e4m3fn.
k_t (MHAOperand): KV cache operand type backing k. Carries the KV layout, dtype, and head geometry.
output_type (DType): Element type of the output tensor output_ptr. bfloat16 when q_type is float8_e4m3fn, else matches q_type.
mask_t (MHAMask): Mask functor type applied to attention scores.
ValidLT (TensorLayout): Compile-time layout of the valid_length_tt tensor.
BM (Int): Number of query rows per block (the M-tile height).
BN (Int): Number of key rows per block (the N-tile width).
BK (Int): Tile size in the depth (head) dimension.
WM (Int): Warp tile height in the M dimension.
WN (Int): Warp tile width in the N dimension.
depth (Int): Total head dimension of Q and K (576 for DeepSeek V2/3).
depth_v (Int): V head dimension; V is derived as K[:, :, :depth_v] (512 for DeepSeek V2/3).
num_heads (Int): Number of query attention heads.
num_threads (Int): Number of threads per CTA, derived from the warp tile geometry.
num_pipeline_stages (Int): Number of software-pipelined MMA stages.
group (Int): GQA group size (num_heads // kv_num_heads); for MLA kv_num_heads == 1 so group == num_heads (defaults to 1).
ragged (Bool): Whether q_ptr is a ragged tensor with batch offsets in valid_length_tt instead of a padded tensor (defaults to False).
_use_valid_length (Bool): Whether to read valid_length_tt for per-batch sequence lengths; False treats all sequences as q_seq_len long (defaults to False).
_is_cache_length_accurate (Bool): Whether max_cache_valid_length is the accurate latest length; False adds seq_len for continuous execution (defaults to False).
decoding_warp_split_k (Bool): Whether to use warp-level split-K reduction within the decode kernel (defaults to False).
q_seq_len (Int): Number of query tokens (S) folded into the MMA M dimension for MTP; 1 for single-token decode (defaults to 1).

Args:

q_ptr (Pointer[Scalar[q_type], MutAnyOrigin, _safe=False]): Pointer to the query tensor with shape [batch, q_seq_len, num_heads, depth] (or ragged equivalent).
k (k_t): KV cache operand. V is derived as K[:, :, :depth_v], so V is not loaded separately.
output_ptr (Pointer[Scalar[output_type], MutAnyOrigin, _safe=False]): Pointer to the output tensor with shape [batch, q_seq_len, num_heads, depth_v] (or ragged equivalent).
exp_sum_ptr (Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False]): Pointer to the per-row softmax exp-sum stats buffer, used for split-K reduction.
qk_max_ptr (Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False]): Pointer to the per-row softmax max stats buffer, used for split-K reduction.
scale (Float32): Softmax scale factor applied to QK^T.
batch_size (Int): Number of sequences in the batch.
num_partitions (Int): Split-K partition count; 1 disables split-K reduction.
max_cache_valid_length (Int): Total number of cached KV entries (the cache context length).
valid_length_tt (TileTensor[DType.uint32, ValidLT, ImmutAnyOrigin]): Per-batch uint32 tensor of valid (pre-padding) sequence lengths. For ragged inputs, carries cumulative row offsets.
mask (mask_t): Mask functor instance applied to attention scores.