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

SplitKPartition

struct SplitKPartition[dtype: DType]

A multi-partition split-K scheme for MHA decoding over long sequences.

Divides the key sequence across num_partitions CTAs. Each CTA writes its partial softmax numerator/denominator to the buffer pointed to by ptr, and a separate reduction kernel merges the results. Over-launches up to max_num_partitions CTAs so the grid shape is stable across varying key lengths (enabling CUDA graph capture).

Parameters

• ​dtype (DType): The accumulator element type used for the partial statistics buffer and the final output.

Fields

• ​ptr (Pointer[Scalar[SplitKPartition[dtype].accum_dtype], MutAnyOrigin]):
• ​num_partitions_value (UInt32):
• ​max_num_partitions_value (UInt32):

Implemented traits

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

comptime members

accum_dtype

comptime accum_dtype = dtype

do_partition

comptime do_partition = True

LSEPointerType

comptime LSEPointerType = NonNullPointer[SplitKPartition[dtype].accum_dtype]

Methods

__init__

def __init__(ptr: Pointer[Scalar[Self.accum_dtype], MutAnyOrigin], num_partitions_value: UInt32, max_num_partitions_value: UInt32) -> Self

num_partitions

def num_partitions(self) -> UInt32

Returns:

UInt32

max_num_partitions

def max_num_partitions(self) -> UInt32

Returns:

UInt32

lse_pointer

def lse_pointer(self) -> Self.LSEPointerType

Returns:

Self.LSEPointerType