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 struct

KVCvt2MmaConsumer

struct KVCvt2MmaConsumer[dtype: DType, config: MLA_SM100_Decode_Config]

Consumes BF16 KV tiles from the convert producer for the MMA pipeline.

Parameters​

  • ​dtype (DType): Element type of the converted KV tiles stored in SMEM.
  • ​config (MLA_SM100_Decode_Config): Decode config supplying num_kv_stages, BN_QK, and q_depth for pipeline stage sizing.

Fields​

  • ​pipe (KVCvt2MmaConsumer[dtype, config].KVPipeType):
  • ​smem (Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

kv_stage_elems​

comptime kv_stage_elems = (config * config)

KVPipeType​

comptime KVPipeType = KVPipelineGeneric[config.num_kv_stages, Int(1), _resolve_warpgroup_size(), Int(2)]

Methods​

__init__​

def __init__(pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), _resolve_warpgroup_size(), Int(2)], smem: Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]) -> Self

stage_base_ptr​

def stage_base_ptr[*, qk_stage: Int = Int(0)](self) -> Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]

Returns:

Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]

stage_index​

def stage_index[*, qk_stage: Int = Int(0)](self) -> UInt32

Returns:

UInt32

wait​

def wait[*, qk_stage: Int = Int(0)](self)

release​

def release[*, qk_stage: Int = Int(0)](mut self, e: Int32)

Was this page helpful?