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 struct

CombineParams

struct CombineParams[output_type: DType, accum_type: DType, num_splits: Int, ragged: Bool = False, warps_per_head: Int = Int(2), has_attn_sink: Bool = False]

Holds the pointers, strides, and shape metadata for the main split-K combine kernel.

Carries 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 by mla_combine_kernel to index the split, batch, sequence, and head dimensions.

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.
• ​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.

Fields

• ​out_accum_split_ptr (Pointer[Scalar[output_type], MutAnyOrigin]):
• ​lse_accum_split_ptr (Pointer[Scalar[accum_type], MutAnyOrigin]):
• ​output_ptr (Pointer[Scalar[output_type], MutAnyOrigin]):
• ​input_row_offsets_ptr (Pointer[UInt32, MutAnyOrigin]):
• ​attn_sink_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin]]):
• ​batch_size (Int):
• ​seq_len (Int):
• ​num_heads (Int):
• ​head_dim (Int):
• ​lse_stride_split (Int):
• ​lse_stride_batch (Int):
• ​lse_stride_seq (Int):
• ​out_accum_stride_split (Int):
• ​out_accum_stride_head (Int):
• ​out_stride_row (Int):

Implemented traits

AnyType, Copyable, Deinitable, DevicePassable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

device_type

comptime device_type = CombineParams[output_type, accum_type, num_splits, ragged, warps_per_head, has_attn_sink]

heads_per_block

comptime heads_per_block = (Int(8) // warps_per_head)

num_threads

comptime num_threads = ((CombineParams[output_type, accum_type, num_splits, ragged, warps_per_head, has_attn_sink].heads_per_block * warps_per_head) * _resolve_warp_size())

Methods

__init__

def __init__(out_accum_split_ptr: Pointer[Scalar[output_type], MutAnyOrigin], lse_accum_split_ptr: Pointer[Scalar[accum_type], MutAnyOrigin], output_ptr: Pointer[Scalar[output_type], MutAnyOrigin], input_row_offsets_ptr: Pointer[UInt32, MutAnyOrigin], attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin]], batch_size: Int, seq_len: Int, num_heads: Int, head_dim: Int) -> Self

get_type_name

static def get_type_name() -> String

Returns:

String

get_device_type_name

static def get_device_type_name() -> String

Returns:

String