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_single_batch

def mla_decoding_single_batch[q_type: DType, k_t: MHAOperand, output_type: DType, mask_t: MHAMask, *, BM: Int, BN: Int, BK: Int, WM: Int, WN: Int, depth: Int, depth_v: Int, num_threads: Int, num_pipeline_stages: Int, decoding_warp_split_k: Bool = False](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, num_keys: Int, num_partitions: Int, mask: mask_t, batch_idx: Int)

Flash attention v2 algorithm.

Computes single-batch MLA attention for one decode step on NVIDIA GPUs: Q @ K^T with online softmax, then P @ V where V is derived as K[:, :depth_v]. Split-K partitions the key range across block_idx.x; block_idx.y selects the query head group.

Parameters:

q_type (DType): Element type of the query tensor q_ptr (inferred). Must be a half-float or float8_e4m3fn.
k_t (MHAOperand): KV cache operand type backing k (inferred). Carries the KV layout, dtype, and head geometry.
output_type (DType): Element type of the output tensor 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).
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_threads (Int): Number of threads per CTA, derived from the warp tile geometry.
num_pipeline_stages (Int): Number of software-pipelined MMA stages.
decoding_warp_split_k (Bool): Whether to use warp-level split-K reduction within the decode kernel (defaults to False). Currently unsupported; must be False.

Args:

q_ptr (Pointer[Scalar[q_type], MutAnyOrigin, _safe=False]): Pointer to the query tensor for this batch, with num_heads rows of depth elements each, indexed by block_idx.y in BM-row blocks.
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 for this batch, with num_heads rows of depth_v elements each, indexed by block_idx.y in BM-row blocks.
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.
num_keys (Int): Total number of cached KV entries for this batch.
num_partitions (Int): Split-K partition count; 1 disables split-K reduction.
mask (mask_t): Mask functor instance applied to attention scores.
batch_idx (Int): Batch index into the paged KV cache.