IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /max/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, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

accum_dtype​

comptime accum_dtype

do_partition​

comptime do_partition

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

get_exp_sum_qk_max_pointer​

def get_exp_sum_qk_max_pointer(self) -> Pointer[Scalar[Self.accum_dtype], MutAnyOrigin, _safe=False]

Returns:

Pointer[Scalar[_Self.accum_dtype], MutAnyOrigin, _safe=False]

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.

Was this page helpful?