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 struct

KVLoad2CvtProducer

struct KVLoad2CvtProducer[dtype: DType, config: MLA_SM100_Decode_Config]

Producer side of the FP8-to-BF16 load-and-convert KV pipeline.

Parameters

• ​dtype (DType): Element type of the FP8 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 (KVLoad2CvtProducer[dtype, config].KVPipeType):
• ​smem_upper_fp8 (Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

bf16_stage_elems

comptime bf16_stage_elems = (config * config)

fp8_stage_stride_elems

comptime fp8_stage_stride_elems = (Int(2) * Int((mul config.BN_QK, config.q_depth)))

KVPipeType

comptime KVPipeType = KVPipelineGeneric[config.num_kv_stages, Int(1), (Int(1) + Int(32) if (xor (eq (xor (lt config.scale_block_size, 1), True), False), True) else Int(0)), (_resolve_warpgroup_size() + Int(2))]

Methods

__init__

def __init__(pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), (Int(1) + Int(32) if (xor (eq (xor (lt config.scale_block_size, 1), True), False), True) else Int(0)), (_resolve_warpgroup_size() + Int(2))], smem_upper_fp8: Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> Self

init

def init(self)

stage_base_ptr

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

Returns:

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

producer_mbar

def producer_mbar[*, qk_stage: Int = Int(0)](self) -> MBarType

Returns:

MBarType

acquire

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

commit_step

def commit_step(mut self)