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

launch_mla_combine_kernel

def launch_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](out_accum_split: TileTensor[output_type, Storage=out_accum_split.Storage, linear_idx_type=out_accum_split.linear_idx_type], lse_accum_split: TileTensor[accum_type, Storage=lse_accum_split.Storage, linear_idx_type=lse_accum_split.linear_idx_type], output: TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type], input_row_offsets_ptr: Pointer[UInt32, MutAnyOrigin, _safe=False], attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]], batch_size: Int, seq_len: Int, num_heads: Int, ctx: DeviceContext)

Launches the main split-K combine kernel on the device with one CTA per (batch, seq, head-block).

Extracts raw device pointers from the input TileTensor arguments, constructs a CombineParams instance, and enqueues mla_combine_kernel with a grid of (batch_size, seq_len, ceildiv(num_heads, heads_per_block)) blocks and PDL launch attributes.

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:

out_accum_split (TileTensor[output_type, Storage=out_accum_split.Storage, linear_idx_type=out_accum_split.linear_idx_type]): The per-split partial output accumulator tensor of shape [num_splits, batch_size, seq_len, num_heads, head_dim].
lse_accum_split (TileTensor[accum_type, Storage=lse_accum_split.Storage, linear_idx_type=lse_accum_split.linear_idx_type]): The per-split LSE accumulator tensor of shape [num_splits, batch_size, seq_len, num_heads].
output (TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): The final combined output tensor of shape [batch_size, seq_len, num_heads, head_dim] (or the ragged-packed equivalent).
input_row_offsets_ptr (Pointer[UInt32, MutAnyOrigin, _safe=False]): Pointer to cumulative token counts per batch; input_row_offsets_ptr[i] is the start token index for batch i. Used only in ragged mode.
attn_sink_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional pointer to per-head attention-sink LSE values of shape [num_heads] in natural log. Used only when has_attn_sink is true.
batch_size (Int): Number of batches in the request.
seq_len (Int): Maximum number of query tokens per batch (the padded grid dimension along the sequence axis).
num_heads (Int): Number of query attention heads.
ctx (DeviceContext): The device context used to enqueue the kernel.