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

grouped_matmul_kernel_sm100

def grouped_matmul_kernel_sm100[a_type: DType, b_type: DType, c_type: DType, static_K: Int, a_tile_rank: Int, a_tile_shape: IndexList[a_tile_rank], a_desc_shape: IndexList[a_tile_rank], b_tile_rank: Int, b_tile_shape: IndexList[b_tile_rank], b_desc_shape: IndexList[b_tile_rank], CLayout: TensorLayout, AOffsetsLayout: TensorLayout, ExpertIdsLayout: TensorLayout, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, c_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, transpose_b: Bool = True, num_threads: 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](a_tma_op: TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape], a_offsets: TileTensor[DType.uint32, AOffsetsLayout, ImmUnsafeAnyOrigin], expert_ids: TileTensor[DType.int32, ExpertIdsLayout, ImmUnsafeAnyOrigin], c: TileTensor[c_type, CLayout, MutAnyOrigin], num_iters: Int)

Computes the SM100 (Blackwell) grouped matmul using TMA async copies into shared memory and tcgen05 tensor-memory MMA accumulation, with an optional elementwise epilogue.

Each block processes one expert tile; A tiles are loaded via TMA into shared memory and accumulated in tensor memory (tcgen05), then drained to registers and stored to global memory (or passed to the elementwise epilogue).

Parameters:

a_type (DType): Element DType of the A activation tensor.
b_type (DType): Element DType of the B weight tensor.
c_type (DType): Element DType of the C output tensor.
static_K (Int): The contraction dimension K, known at compile time.
a_tile_rank (Int): Rank of the A TMA descriptor's tile shape.
a_tile_shape (IndexList[a_tile_rank]): Per-copy tile shape for A TMA loads from global to shared memory.
a_desc_shape (IndexList[a_tile_rank]): Global-tensor descriptor shape backing the A TMA tile.
b_tile_rank (Int): Rank of the B TMA descriptor's tile shape.
b_tile_shape (IndexList[b_tile_rank]): Per-copy tile shape for B TMA loads from global to shared memory.
b_desc_shape (IndexList[b_tile_rank]): Global-tensor descriptor shape backing the B TMA tile.
CLayout (TensorLayout): TensorLayout of the output tensor c.
AOffsetsLayout (TensorLayout): TensorLayout of the a_offsets tensor.
ExpertIdsLayout (TensorLayout): TensorLayout of the expert_ids tensor.
block_tile_shape (IndexList[Int(3)]): Per-block tile dimensions (BM, BN, BK)` for the M, N, and K axes.
mma_shape (IndexList[Int(3)]): tcgen05 MMA instruction shape (MMA_M, MMA_N, MMA_K).
a_swizzle (TensorMapSwizzle): TMA swizzle mode for A shared-memory loads (defaults to SWIZZLE_128B).
b_swizzle (TensorMapSwizzle): TMA swizzle mode for B shared-memory loads (defaults to SWIZZLE_128B).
c_swizzle (TensorMapSwizzle): Swizzle mode for C output stores (defaults to SWIZZLE_NONE).
transpose_b (Bool): Whether B is stored in transposed layout (defaults to True).
num_threads (Int): Number of threads per block, either 128 or 256 (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 elementwise epilogue applied to each output element (defaults to None).