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_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], 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, 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]): 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]): 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]): 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.
• ​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.