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
DecodeSM100QKTTS
struct DecodeSM100QKTTS[operand_type: DType, accum_type: DType, *, config: MLA_SM100_Decode_Config]
Tensor accumulator for the QK^T MMA with Q in TMEM and K in SMEM.
Parameters
- operand_type (
DType): Element type of the Q (in TMEM) and K (in SMEM) operands. - 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 = MLA_SM100_Decode_Config.MMA_K
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(16))
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_k_block
static def descriptor_k_block(kv_smem: Pointer[Scalar[operand_type], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
mma
static def mma[*, stage_idx: Int = Int(0)](a: UInt32, b: MMASmemDescriptorPair, c: UInt32, *, c_scale: UInt32, elect: Int32)