For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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], k: k_t, output_ptr: Pointer[Scalar[output_type], MutAnyOrigin], exp_sum_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin], qk_max_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin], scale: Float32, batch_size_dev: Int32, num_partitions_dev: Int32, max_cache_valid_length_dev: Int32, 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]): 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]): 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]): 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]): 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_dev (Int32): Number of sequences in the batch.
• ​num_partitions_dev (Int32): Split-K partition count; 1 disables split-K reduction.
• ​max_cache_valid_length_dev (Int32): 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.