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

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], output_ptr: Pointer[Scalar[output_type], MutAnyOrigin], exp_sum_ptr: Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin], qk_max_ptr: Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin], batch_size: Int32, num_partitions: Int32)

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]): Pointer to partial output buffer [num_partitions, batch, heads, depth].
• ​output_ptr (Pointer[Scalar[output_type], MutAnyOrigin]): Pointer to final output buffer [batch, heads, depth].
• ​exp_sum_ptr (Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin]): Pointer to partial exp-sum buffer [num_partitions, batch, heads].
• ​qk_max_ptr (Pointer[Scalar[get_accum_type[output_type]()], ImmutAnyOrigin]): Pointer to partial softmax-max buffer [num_partitions, batch, heads].
• ​batch_size (Int32): Number of sequences in the batch.
• ​num_partitions (Int32): Number of split-K partitions to reduce.