Skip to main content
Log in

Mojo struct

SplitKPartition

@register_passable(trivial) struct SplitKPartition[dtype: DType]

Aliases

  • do_partition = True:
  • accum_dtype = dtype:

Fields

  • ptr (UnsafePointer[SIMD[dtype, 1]]):
  • num_partitions_value (SIMD[uint32, 1]):

Implemented traits

AnyType, Copyable, ExplicitlyCopyable, MHAPartitionScheme, Movable, UnknownDestructibility

Methods

__init__

__init__(ptr: UnsafePointer[SIMD[dtype, 1]], num_partitions_value: SIMD[uint32, 1]) -> Self

num_partitions

num_partitions(self) -> SIMD[uint32, 1]

get_exp_sum_qk_max_pointer

get_exp_sum_qk_max_pointer(self) -> UnsafePointer[SIMD[dtype, 1]]