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

gemm_mma_cpasync_kernel

def gemm_mma_cpasync_kernel[c_type: DType, a_type: DType, b_type: DType, c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, *, tile_m: Int = Int(16), tile_n: Int = Int(8), tile_k: Int = Int(128), stage_cnt: Int = Int(2), accum_type: DType = DType.float32, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, pdl_level: PDLLevel = PDLLevel(), swapAB: Bool = False](output: TileTensor[c_type, c_layout, MutAnyOrigin], act: TileTensor[a_type, a_layout, ImmutAnyOrigin], weight: TileTensor[b_type, b_layout, ImmutAnyOrigin], gemm_m: Int, gemm_k: Int, gemm_n: Int, batch_size: Int)

SM100 (B200) warp-specialized GEMM kernel using cp.async and m16n8k16 MMA instructions.

Implements a multi-stage software-pipelining loop: producer warps (0-3) prefetch activation and weight tiles from global memory into shared memory via cp.async, and consumer warps (4+) issue m16n8k16 MMA operations and write results to the output. Requires B200 (sm_100x) hardware.

Parameters:

c_type (DType): Output element type.
a_type (DType): Activation (A) element type.
b_type (DType): Weight (B) element type.
c_layout (TensorLayout): TensorLayout for the output.
a_layout (TensorLayout): TensorLayout for activations.
b_layout (TensorLayout): TensorLayout for weights.
tile_m (Int): Tile size along M; must be 16 for m16n8k16 MMA.
tile_n (Int): Tile size along N per warp.
tile_k (Int): Tile size along K; determines shared memory allocation.
stage_cnt (Int): Number of pipeline stages.
accum_type (DType): Accumulator precision; must be float32.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue applied to each output element.
pdl_level (PDLLevel): Programmatic dependent launch level.
swapAB (Bool): When True, swaps the A and B operand roles.

Args:

output (TileTensor[c_type, c_layout, MutAnyOrigin]): Rank-3 output TileTensor of shape (batch, gemm_m, gemm_n).
act (TileTensor[a_type, a_layout, ImmutAnyOrigin]): Rank-3 activation TileTensor of shape (batch, gemm_m, gemm_k).
weight (TileTensor[b_type, b_layout, ImmutAnyOrigin]): Rank-3 weight TileTensor of shape (batch, gemm_n, gemm_k).
gemm_m (Int): M dimension of the GEMM.
gemm_k (Int): K (reduction) dimension of the GEMM.
gemm_n (Int): N dimension of the GEMM.
batch_size (Int): Number of independent GEMM problems in the batch.