IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /get-started.md).

Mojo struct

DecodeSM100QKTSS_FP8

struct DecodeSM100QKTSS_FP8[operand_type: DType, accum_type: DType, *, config: MLA_SM100_Decode_Config]

Tensor accumulator for the native FP8 QK^T MMA with both Q and K in FP8 SMEM.

Parameters

  • operand_type (DType): FP8 element type of the Q and K operands in SMEM.
  • accum_type (DType): Accumulator dtype used for the QK^T MMA result in TMEM.
  • config (MLA_SM100_Decode_Config): Decode config supplying MMA tile dimensions and swizzle modes.

Implemented traits

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

comptime members

BK

comptime BK = config.BK_QK

MMA_K

comptime MMA_K = 32

MMA_M

comptime MMA_M = config.MMA_M

MMA_N

comptime MMA_N = config.MMA_QK_N

num_k_mmas

comptime num_k_mmas = (config // Int(32))

operand_size

comptime operand_size = size_of[operand_type]()

UMMAInstDesc

comptime UMMAInstDesc = UMMAInsDescriptor.create[accum_type, operand_type, operand_type, Index[Int, Int, dtype=DType.uint32](config, config)]()

Methods

descriptor_q_block

static def descriptor_q_block(q_smem: Pointer[Scalar[operand_type], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

descriptor_k_block

static def descriptor_k_block(kv_smem: Pointer[Scalar[operand_type], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

mma

static def mma[*, stage_idx: Int = Int(0)](a: MMASmemDescriptorPair, b: MMASmemDescriptorPair, c: UInt32, *, c_scale: UInt32, elect: Int32)