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
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
max_num_partitions
get_exp_sum_qk_max_pointer
def get_exp_sum_qk_max_pointer(self) -> Pointer[Scalar[Self.accum_dtype], MutAnyOrigin]
Returns:
Pointer[Scalar[_Self.accum_dtype], MutAnyOrigin]
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.