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 trait

MHAPartitionScheme

Trait describing how the key-value sequence is partitioned for split-K decoding.

Implementations either skip partitioning entirely (NoPartition) or divide the key sequence across multiple CTAs and accumulate partial softmax statistics in a separate reduction pass (SplitKPartition). The do_partition compile-time flag lets the compiler eliminate the reduction kernel when no partitioning is needed.

Implemented traits

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

comptime members

accum_dtype

comptime accum_dtype

do_partition

comptime do_partition

LSEPointerType

comptime LSEPointerType

Required methods

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