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

def gemm_mma_cpasync[pdl_level: PDLLevel = PDLLevel(), tile_k: Int = Int(128), elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, swapAB: Bool = False](c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], act: TileTensor[Storage=act.Storage, address_space=act.address_space, linear_idx_type=act.linear_idx_type], weight: TileTensor[Storage=weight.Storage, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type], gemm_m: Int, gemm_k: Int, gemm_n: Int, batch_size: Int, ctx: DeviceContext)

Launch the batched GEMM tensor-core kernel.

C[gemm_m, gemm_n] = act[gemm_m, K] x weight[gemm_n, K]^T.

The caller always passes act/weight/c with the same shapes regardless of swapAB, and c is always written as a row-major [gemm_m, gemm_n] ([M, N]) buffer. swapAB only changes the internal tiling: when True the weight is fed to the A (free-streaming) operand slot and the activation to the B slot (so the grid tiles N by tile_m and M by tile_n), and the epilogue transposes the store back into the row-major [M, N] buffer. This makes the large weight the producer-independent operand for PDL overlap at small M (decode), while keeping the output layout identical.

Parameters:

pdl_level (PDLLevel): Programmatic dependent launch level for PDL barriers.
tile_k (Int): K-dimension tile size for the MMA kernel (defaults to 128).
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.
swapAB (Bool): When True, feeds the weight to the A operand slot and the activation to the B slot for PDL overlap at small M.

Args:

c (TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output, shape (gemm_m, gemm_n) or (batch, gemm_m, gemm_n), always row-major.
act (TileTensor[Storage=act.Storage, address_space=act.address_space, linear_idx_type=act.linear_idx_type]): Activation, shape (gemm_m, gemm_k) or (batch, gemm_m, gemm_k).
weight (TileTensor[Storage=weight.Storage, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type]): Weight, shape (gemm_n, gemm_k) or (batch, gemm_n, gemm_k).
gemm_m (Int): Activation rows (output rows, M).
gemm_k (Int): Reduction dimension.
gemm_n (Int): Weight rows (output cols, N).
batch_size (Int): Batch size; ignored for 2D inputs (treated as 1).
ctx (DeviceContext): GPU device context.