IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

__init__

def __init__(out self, *, copy: Self)

Create a new instance of the value by copying an existing one.

Args:

  • copy (_Self): The value to copy.

Returns:

_Self

def __init__(out self, *, deinit move: Self)

Create a new instance of the value by moving the value of another.

Args:

  • move (_Self): The value to move.

Returns:

_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

Provided methods

copy

def copy(self) -> Self

Explicitly construct a copy of self, a convenience method for Self(copy=self) when the type is inconvenient to write out.

Overriding this method is not allowed.

Returns:

_Self: A copy of this value.