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 function

bulk_mma

def bulk_mma[kind: UMMAKind, //, layout_a: Layout, layout_b: Layout, *, num_k_mmas: Int, mma_k: Int, operand_size: Int, cta_group: Int = Int(1)](idesc: UMMAInsDescriptor[kind], a: MMASmemDescriptorPair, b: MMASmemDescriptorPair, c_tmem: UInt32, c_scale: UInt32, elect: Int32)

Issues a full-tile SS (both operands in SMEM) non-warp-specialized tcgen05.mma contraction.

Parameters:

kind (UMMAKind): UMMAKind selecting the tcgen05.mma instruction variant.
layout_a (Layout): SMEM layout of the A operand tile, used to compute per-K-block A descriptor offsets.
layout_b (Layout): SMEM layout of the B operand tile, used to compute per-K-block B descriptor offsets.
num_k_mmas (Int): Number of mma_k-sized K-dimension blocks to contract over.
mma_k (Int): K-dimension tile size per MMA block, in elements.
operand_size (Int): Size in bytes of the A and B operand elements.
cta_group (Int): Number of cooperating CTAs, 1 or 2 (defaults to 1).

Args:

idesc (UMMAInsDescriptor[kind]): UMMA instruction descriptor encoding the accumulator and operand dtypes and the output tile shape.
a (MMASmemDescriptorPair): SMEM descriptor pair for the A operand.
b (MMASmemDescriptorPair): SMEM descriptor pair for the B operand.
c_tmem (UInt32): TMEM base address of the output accumulator C.
c_scale (UInt32): Accumulator init/accumulate scale; nonzero on the first block to initialize the accumulator, zero to accumulate.
elect (Int32): elect() result selecting the single thread that issues the MMA.

def bulk_mma[kind: UMMAKind, //, layout_b: Layout, *, mma_k: Int, num_k_mmas: Int, operand_size: Int, cta_group: Int = Int(1)](idesc: UMMAInsDescriptor[kind], a: UInt32, b: MMASmemDescriptorPair, c_tmem: UInt32, c_scale: UInt32, elect: Int32)

Issues a full-tile TS (A in TMEM, B in SMEM) non-warp-specialized tcgen05.mma contraction.

Parameters:

kind (UMMAKind): UMMAKind selecting the tcgen05.mma instruction variant.
layout_b (Layout): SMEM layout of the B operand tile, used to compute per-K-block B descriptor offsets.
mma_k (Int): K-dimension tile size per MMA block, in elements.
num_k_mmas (Int): Number of mma_k-sized K-dimension blocks to contract over.
operand_size (Int): Size in bytes of the A and B operand elements.
cta_group (Int): Number of cooperating CTAs, 1 or 2 (defaults to 1).

Args:

idesc (UMMAInsDescriptor[kind]): UMMA instruction descriptor encoding the accumulator and operand dtypes and the output tile shape.
a (UInt32): TMEM base address of the A operand.
b (MMASmemDescriptorPair): SMEM descriptor pair for the B operand.
c_tmem (UInt32): TMEM base address of the output accumulator C.
c_scale (UInt32): Accumulator init/accumulate scale; nonzero on the first block to initialize the accumulator, zero to accumulate.
elect (Int32): elect() result selecting the single thread that issues the MMA.