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 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, _safe=False]): - num_partitions_value (
UInt32): - max_num_partitions_value (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
MHAPartitionScheme,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
accum_dtype
comptime accum_dtype = dtype
do_partition
comptime do_partition = True
Methods
__init__
def __init__(ptr: Pointer[Scalar[Self.accum_dtype], MutAnyOrigin, _safe=False], num_partitions_value: UInt32, max_num_partitions_value: UInt32) -> 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, _safe=False]
Returns:
Pointer[Scalar[Self.accum_dtype], MutAnyOrigin, _safe=False]
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!