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_splitk_reduce

def mla_splitk_reduce[intermediate_type: DType, output_type: DType, ValidLT: TensorLayout, depth: Int, num_heads: Int, D_TILES: Int, W_PARTS: Int, MAX_PARTITIONS: Int, use_exp2: Bool = False, q_seq_len: Int = Int(1), ragged: Bool = False](intermediate_ptr: Pointer[Scalar[intermediate_type], ImmutAnyOrigin, _safe=False], output_ptr: Pointer[Scalar[output_type], MutAnyOrigin, _safe=False], exp_sum_ptr: Pointer[Scalar[get_accum_type[output_type]()], MutAnyOrigin, _safe=False], qk_max_ptr: Pointer[Scalar[get_accum_type[output_type]()], MutAnyOrigin, _safe=False], batch_size: Int, num_partitions: Int, valid_length_tt: TileTensor[DType.uint32, ValidLT, ImmutAnyOrigin])

Combines per-partition MLA decode partial outputs into the final result.

Performs the split-K reduction for MLA decoding by merging per-partition intermediate outputs and softmax statistics (exp-sum and qk-max) across the partition axis using W_PARTS warps per CTA and D_TILES CTAs over the depth axis, then stores the rescaled output to global memory.

Parameters:

intermediate_type (DType): Element type of the per-partition intermediate output buffer.
output_type (DType): Element type of the final reduced output.
ValidLT (TensorLayout): Layout type of the valid_length_tt tensor (inferred).
depth (Int): Per-head attention head dimension. Must be positive and divisible by D_TILES * WARP_SIZE; 576 for DeepSeek V2/3.
num_heads (Int): Number of query attention heads. The folded query-row count is num_heads * q_seq_len.
D_TILES (Int): Number of CTAs the kernel launches along the depth axis; each CTA owns depth / D_TILES elements.
W_PARTS (Int): Number of warps per CTA, each handling MAX_PARTITIONS / W_PARTS partitions.
MAX_PARTITIONS (Int): Upper bound on the partition count this specialization handles; must be a multiple of WARP_SIZE and W_PARTS.
use_exp2 (Bool): Whether the softmax rescale uses exp2 instead of exp (defaults to False). AMD uses exp2.
q_seq_len (Int): Query tokens per batch (the MTP token-fold ceiling); the MMA M dimension is num_heads * q_seq_len (defaults to 1).
ragged (Bool): Whether the final store is remapped through valid_length_tt to the ragged [total_q_tokens, num_heads, depth] layout (defaults to False).

Args:

intermediate_ptr (Pointer[Scalar[intermediate_type], ImmutAnyOrigin, _safe=False]): Pointer to the per-partition intermediate output with shape [num_partitions, batch_size, num_rows, depth] where num_rows = num_heads * q_seq_len.
output_ptr (Pointer[Scalar[output_type], MutAnyOrigin, _safe=False]): Pointer to the final output buffer with shape [batch_size * num_rows, depth].
exp_sum_ptr (Pointer[Scalar[get_accum_type[output_type]()], MutAnyOrigin, _safe=False]): Pointer to per-partition exp-sum statistics with shape [num_partitions, batch_size, num_rows].
qk_max_ptr (Pointer[Scalar[get_accum_type[output_type]()], MutAnyOrigin, _safe=False]): Pointer to per-partition qk-max statistics with shape [num_partitions, batch_size, num_rows].
batch_size (Int): Number of sequences in the batch.
num_partitions (Int): Number of split-K partitions actually used. Must be positive and at most MAX_PARTITIONS.
valid_length_tt (TileTensor[DType.uint32, ValidLT, ImmutAnyOrigin]): uint32 tensor of input row offsets of shape [batch_size + 1]. Read only on the ragged and q_seq_len > 1 path to remap the final store; a zero-length placeholder otherwise.