Skip to main content

Mojo struct

KVLoad2CvtConsumer

struct KVLoad2CvtConsumer[dtype: DType, config: MLA_SM100_Decode_Config]

Fields​

  • ​pipe (KVLoad2CvtConsumer[dtype, config].PipeT):
  • ​smem_upper_fp8 (UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

bf16_stage_elems​

comptime bf16_stage_elems = (config * config)

fp8_stage_stride_elems​

comptime fp8_stage_stride_elems = (2 * KVLoad2CvtConsumer[dtype, config].bf16_stage_elems)

PipeT​

comptime PipeT = KVPipelineGeneric[config.num_kv_stages, 1, (1 + 32 if (config > 0) else 0), (WARPGROUP_SIZE + 2)]

Methods​

__init__​

__init__(pipe: KVPipelineGeneric[config.num_kv_stages, 1, (1 + 32 if (config > 0) else 0), (WARPGROUP_SIZE + 2)], smem_upper_fp8: UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> Self

stage_base_ptr​

stage_base_ptr(self) -> UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]

Returns:

UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]

wait​

wait(self)

release_all​

release_all(mut self)