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

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

AnyType, Deinitable, Movable

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]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

tmem_descriptor_q

static def tmem_descriptor_q(q_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

descriptor_k_p0

static def descriptor_k_p0(k_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

descriptor_k_p1

static def descriptor_k_p1(k_smem: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair