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

mha_splitk_reduce

def mha_splitk_reduce[intermediate_type: DType, output_type: DType, depth: Int, num_heads: Int, num_threads: Int, use_exp2: 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]()], ImmutAnyOrigin, _safe=False], qk_max_ptr: Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin, _safe=False], batch_size: Int, num_partitions: Int)

Single-warp reduction kernel that merges split-K partial attention outputs.

Reads num_partitions partial attention outputs together with their running exp_sum (denominator) and qk_max statistics, re-weights each partial output by its softmax scale relative to the global maximum, sums them, and writes the normalised result to output_ptr. Must be launched with exactly one warp (WARP_SIZE threads) per output row.

Parameters:

intermediate_type (DType): Element type of the partial output buffer written by the split-K decode kernel.
output_type (DType): Element type of the final attention output.
depth (Int): Attention head depth.
num_heads (Int): Number of query heads.
num_threads (Int): Must equal WARP_SIZE.
use_exp2 (Bool): True to use base-2 exponentiation for numerics matching FA3 kernels that fuse scale * log2e.

Args:

intermediate_ptr (Pointer[Scalar[intermediate_type], ImmutAnyOrigin, _safe=False]): Pointer to partial output buffer [num_partitions, batch, heads, depth].
output_ptr (Pointer[Scalar[output_type], MutAnyOrigin, _safe=False]): Pointer to final output buffer [batch, heads, depth].
exp_sum_ptr (Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin, _safe=False]): Pointer to partial exp-sum buffer [num_partitions, batch, heads].
qk_max_ptr (Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin, _safe=False]): Pointer to partial softmax-max buffer [num_partitions, batch, heads].
batch_size (Int): Number of sequences in the batch.
num_partitions (Int): Number of split-K partitions to reduce.