Mojo struct
QKMMAOp
struct QKMMAOp[dtype: DType, accum_dtype: DType, config: MLASparseConfig[config.qkv_dtype]]
Implemented traits
AnyType,
ImplicitlyDestructible
comptime members
NUM_TS_STAGES
comptime NUM_TS_STAGES = 3
SSMMAType
comptime SSMMAType = SM100TensorAccumulatorSS[dtype, accum_dtype, config.num_q_heads, 128, 192, cta_group=2]
TSMMAType
comptime TSMMAType = SM100TensorAccumulatorTS[dtype, accum_dtype, config.num_q_heads, 128, 384, cta_group=2, num_stages=3]
Methods
smem_descriptor_q
static smem_descriptor_q(q_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
tmem_descriptor_q
static tmem_descriptor_q(q_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
descriptor_k_p0
static descriptor_k_p0(k_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
descriptor_k_p1
static descriptor_k_p1(k_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!