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
MmaOpSM100_SS
struct MmaOpSM100_SS[c_type: DType, a_type: DType, b_type: DType, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], /, *, accum_type: DType = DType.float32, cta_group: Int = Int(1), cluster_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(1), Int(1), Int(1)), a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, transpose_b: Bool = False]
SM100 (Blackwell) warp-specialized MMA operation for standard (non-block-scaled) GEMM.
Encapsulates a UMMA instruction descriptor and multicast mask for issuing asynchronous tensor-core MMA operations from shared memory to tensor memory (TMEM) on NVIDIA SM100 GPUs. Supports optional CTA clustering for multicast.
Parameters
- c_type (
DType): Output accumulator element type. - a_type (
DType): A operand element type. - b_type (
DType): B operand element type; must matcha_type. - block_tile_shape (
IndexList[Int(3)]): (BM, BN, BK) shape of the shared-memory tile processed per CTA. - mma_shape (
IndexList[Int(3)]): (MMA_M, MMA_N, MMA_K) shape of a single UMMA instruction. - accum_type (
DType): Accumulator precision; defaults to float32. - cta_group (
Int): Number of CTAs collaborating on one tile (1 or 2). - cluster_shape (
IndexList[Int(3)]): CTA cluster shape for multicast; defaults to (1, 1, 1). - a_swizzle (
TensorMapSwizzle): Swizzle mode for the A operand shared-memory layout. - b_swizzle (
TensorMapSwizzle): Swizzle mode for the B operand shared-memory layout. - transpose_b (
Bool): Must be True; SM100 UMMA always uses transposed B.
Fields
- idesc (
UMMAInsDescriptor[MmaOpSM100_SS._get_umma_kind[a_type]()]): - mask (
UInt16):
Implemented traits
AnyType,
Copyable,
Defaultable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
def __init__() -> Self
make_a_desc
def make_a_desc(self, a: TileTensor[Storage=a.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a.linear_idx_type]) -> MMASmemDescriptor
Build the K-major MMA descriptor for an A operand tile.
Exposed so callers that issue MMAs over several adjacent SMEM tiles (e.g. a k-group loop) can build the descriptor once and advance it by a comptime byte stride per tile, rather than rebuilding the full descriptor (base-pointer mask + bitfield inserts) for each tile.
Args:
- a (
TileTensor[Storage=a.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a.linear_idx_type]): A operand tile in shared memory.
Returns:
MMASmemDescriptor: The K-major MMA shared-memory descriptor for a.
make_b_desc
def make_b_desc(self, b: TileTensor[Storage=b.Storage, address_space=AddressSpace.SHARED, linear_idx_type=b.linear_idx_type]) -> MMASmemDescriptor
Build the K-major MMA descriptor for a B operand tile.
See make_a_desc for why this is exposed.
Args:
- b (
TileTensor[Storage=b.Storage, address_space=AddressSpace.SHARED, linear_idx_type=b.linear_idx_type]): B operand tile in shared memory.
Returns:
MMASmemDescriptor: The K-major MMA shared-memory descriptor for b.
mma
def mma(self, a: TileTensor[Storage=a.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a.linear_idx_type], b: TileTensor[Storage=b.Storage, address_space=AddressSpace.SHARED, linear_idx_type=b.linear_idx_type], c_tmem: UInt32, init_c: Bool)
Issue MMA operations over K tiles from shared memory to TMEM.
Args:
- a (
TileTensor[Storage=a.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a.linear_idx_type]): A operand tile in shared memory. - b (
TileTensor[Storage=b.Storage, address_space=AddressSpace.SHARED, linear_idx_type=b.linear_idx_type]): B operand tile in shared memory. - c_tmem (
UInt32): TMEM address for the accumulator. - init_c (
Bool): When True, zero-initialize the accumulator on the first K slice instead of accumulating.
mma_from_desc
def mma_from_desc(self, a_desc: MMASmemDescriptor, b_desc: MMASmemDescriptor, c_tmem: UInt32, init_c: Bool)
Issue MMA operations over K tiles from precomputed SMEM descriptors.
Identical to mma, but takes already-built A/B descriptors. This lets
a k-group caller build the descriptor base once and advance it by a
comptime byte stride per tile (the SM100 SMEM-descriptor base address
adds linearly), avoiding a redundant full descriptor rebuild per tile.
Args:
- a_desc (
MMASmemDescriptor): Precomputed K-major descriptor for the A operand tile. - b_desc (
MMASmemDescriptor): Precomputed K-major descriptor for the B operand tile. - c_tmem (
UInt32): TMEM address for the accumulator. - init_c (
Bool): When True, zero-initialize the accumulator on the first K slice instead of accumulating.
commit
def commit(self, ptr_mbar: Pointer[address_space=AddressSpace.SHARED, _safe=False])
wait
def wait(self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!