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_, config.cta_group_]]
Implemented traits
comptime members
NUM_TS_STAGES
comptime NUM_TS_STAGES = 3
SSMMAType
comptime SSMMAType = SM100TensorAccumulator[dtype, accum_dtype, config.padded_num_q_heads, config.B_TOPK, config.q_smem_depth, a_tmem=False, cta_group=config.cta_group]
TSMMAType
comptime TSMMAType = SM100TensorAccumulator[dtype, accum_dtype, config.padded_num_q_heads, config.B_TOPK, config.q_tmem_depth, a_tmem=True, cta_group=config.cta_group, num_stages=Int(3)]
Methods
smem_descriptor_q
static def smem_descriptor_q(q_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED, _safe=False]) -> MMASmemDescriptorPair
Returns:
tmem_descriptor_q
static def tmem_descriptor_q(q_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED, _safe=False]) -> MMASmemDescriptorPair
Returns:
descriptor_k_p0
static def descriptor_k_p0(k_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED, _safe=False]) -> MMASmemDescriptorPair
Returns:
descriptor_k_p1
static def descriptor_k_p1(k_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED, _safe=False]) -> MMASmemDescriptorPair
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!