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_combine_kernel

def mla_combine_kernel[output_type: DType, accum_type: DType, head_dim: Int, num_splits: Int, ragged: Bool = False, warps_per_head: Int = Int(2), has_attn_sink: Bool = False](params: CombineParams[output_type, accum_type, num_splits, ragged, warps_per_head, has_attn_sink])

Combines partial split-K attention outputs into a single decoded result using LSE-based weighting.

Loads each split's partial output and LSE value, computes a global LSE via warp reduction, derives per-split scale factors as exp2(lse_i - global_lse), and accumulates the weighted sum of partial outputs into the final output tensor. Supports ragged batch layouts and optional attention-sink correction.

Parameters:

output_type (DType): The element type of the per-split partial output accumulator and the final combined output tensor.
accum_type (DType): The element type of the per-split LSE accumulator values.
head_dim (Int): The dimension of each attention head. Must be divisible by WARP_SIZE * warps_per_head.
num_splits (Int): The compile-time number of KV cache splits to combine. Used to unroll the per-split accumulation loop.
ragged (Bool): Whether batches have variable-length sequences in a packed layout (defaults to False). When true, input_row_offsets_ptr selects each batch's token range.
warps_per_head (Int): Number of warps assigned to each attention head (defaults to 2). Must divide 8; controls vector load width and threads per block.
has_attn_sink (Bool): Whether to apply attention-sink correction to the global LSE (defaults to False). When true, attn_sink_ptr must be provided.

Args:

params (CombineParams[output_type, accum_type, num_splits, ragged, warps_per_head, has_attn_sink]): The CombineParams instance carrying the per-split partial output and LSE accumulators, the final output tensor, optional ragged row offsets and attention-sink pointer, and the derived strides used to index the split, batch, sequence, and head dimensions.