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
QKMMAOp
struct QKMMAOp[dtype: DType, accum_dtype: DType, config: MLASparseConfig[config.qkv_dtype, config.b_topk_, config.num_mbars_, config.q_smem_depth_, config.q_tmem_depth_]]
Implemented traits
comptime members
NUM_TS_STAGES
comptime NUM_TS_STAGES = 3
SSMMAType
comptime SSMMAType = SM100TensorAccumulator[dtype, accum_dtype, config.num_q_heads, config.B_TOPK, config.q_smem_depth, a_tmem=False, cta_group=2]
TSMMAType
comptime TSMMAType = SM100TensorAccumulator[dtype, accum_dtype, config.num_q_heads, config.B_TOPK, config.q_tmem_depth, a_tmem=True, cta_group=2, num_stages=3]
Methods
smem_descriptor_q
static def smem_descriptor_q(q_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
tmem_descriptor_q
static def tmem_descriptor_q(q_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
descriptor_k_p0
static def descriptor_k_p0(k_smem: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
descriptor_k_p1
static def 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!