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

blackwell_tma_umma_warp_specialized_kernel

def blackwell_tma_umma_warp_specialized_kernel[a_type: DType, b_type: DType, c_type: DType, expert_m: 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], c_tile_rank: Int, c_tile_shape_param: IndexList[c_tile_rank], c_desc_shape: IndexList[c_tile_rank], block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], cluster_shape: StaticTuple[Int32, Int(3)], num_pipeline_stages: Int, num_accum_pipeline_stages: Int, num_output_stages: Int = Int(2), output_tile_shape: IndexList[Int(2)] = Index[Int, Int](Int(128), Int(32)), transpose_b: Bool = True, a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, c_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, cta_group: Int = Int(2), elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, a_plane_splits: IndexList[Int(2)] = Index[Int, Int](Int(0), Int(0)), transpose_c: Bool = False, use_tma: Bool = True, K_actual: Int = Int(0), a_gmem_layout: Layout = Layout.row_major(Int(1), Int(1)), b_gmem_layout: Layout = Layout.row_major(Int(1), Int(1))](expert_usage_stats: Pointer[UInt32, ImmutAnyOrigin, _safe=False], a_tma_op: TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape], expert_ids: Pointer[Int32, ImmutAnyOrigin, _safe=False], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape], b_offsets: Pointer[UInt32, ImmutAnyOrigin, _safe=False], c_tma_op: TMATensorTile[c_type, c_tile_rank, c_tile_shape_param, c_desc_shape], c_ptr: Pointer[Scalar[c_type], MutAnyOrigin, _safe=False], mnk: StaticTuple[UInt32, Int(3)], a_gmem: LayoutTensor[a_type, a_gmem_layout, ImmutAnyOrigin], b_gmem: LayoutTensor[b_type, b_gmem_layout, ImmutAnyOrigin])

Implements the warp-specialized persistent grouped GEMM kernel for SM100.

Splits the CTA warps into main-load, MMA, and epilogue roles that cooperate through mbarriers and a tile scheduler to iterate over grouped GEMM work tiles, streaming A and B through shared memory via TMA and accumulating into tensor memory before storing C back to global memory.

Parameters:

a_type (DType): Element type of the A operand matrix.
b_type (DType): Element type of the B operand matrix.
c_type (DType): Element type of the C output matrix.
expert_m (Int): Static M dimension of each expert's local slice, used as the C stride and row bound for contiguous row-major C output.
a_tile_rank (Int): Rank of the A TMA tile descriptor.
a_tile_shape (IndexList[a_tile_rank]): Element shape of one A TMA tile.
a_desc_shape (IndexList[a_tile_rank]): Descriptor shape of the A TMA tile, used to compute per- load element counts and row stride.
b_tile_rank (Int): Rank of the B TMA tile descriptor.
b_tile_shape (IndexList[b_tile_rank]): Element shape of one B TMA tile.
b_desc_shape (IndexList[b_tile_rank]): Descriptor shape of the B TMA tile, used to compute per- load element counts and row stride.
c_tile_rank (Int): Rank of the C TMA tile descriptor.
c_tile_shape_param (IndexList[c_tile_rank]): Element shape of one C TMA tile.
c_desc_shape (IndexList[c_tile_rank]): Descriptor shape of the C TMA tile.
block_tile_shape (IndexList[Int(3)]): Block tile shape [BM, BN, BK] partitioning the GEMM into work tiles.
mma_shape (IndexList[Int(3)]): MMA instruction shape [MMA_M, MMA_N, MMA_K].
cluster_shape (StaticTuple[Int32, Int(3)]): Thread block cluster shape [CLUSTER_M, CLUSTER_N, CLUSTER_D].
num_pipeline_stages (Int): Number of TMA producer and consumer pipeline stages used to double-buffer A and B shared-memory tiles.
num_accum_pipeline_stages (Int): Number of accumulator pipeline stages used to double-buffer TMEM-to-global-memory C stores.
num_output_stages (Int): Number of output shared-memory buffer stages for C (defaults to 2).
output_tile_shape (IndexList[Int(2)]): Output tile shape [OUT_M, OUT_N] for one C shared- memory stage (defaults to Index(128, 32)).
transpose_b (Bool): Whether B is stored transposed in global memory (defaults to True).
a_swizzle (TensorMapSwizzle): TMA swizzle mode applied to A shared-memory tiles (defaults to SWIZZLE_128B).
b_swizzle (TensorMapSwizzle): TMA swizzle mode applied to B shared-memory tiles (defaults to SWIZZLE_128B).
c_swizzle (TensorMapSwizzle): TMA swizzle mode applied to C shared-memory tiles (defaults to SWIZZLE_128B).
cta_group (Int): Number of CTAs cooperating per MMA along the M dimension (defaults to 2).
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 stored C fragments (defaults to None).
a_plane_splits (IndexList[Int(2)]): Per-plane split sizes for fused LoRA QKV A-plane row offsetting; (0, 0) disables it (defaults to (0, 0)).
transpose_c (Bool): Whether to transpose the C output tile in shared memory before the TMA store (defaults to False).
use_tma (Bool): Whether to use TMA for A and B loads; when False, use the CUDA-core fallback (defaults to True).
K_actual (Int): Actual K dimension in elements for the CUDA-core fallback when use_tma is False (defaults to 0).
a_gmem_layout (Layout): Layout of the A global-memory tensor, used only by the CUDA-core fallback (defaults to Layout.row_major(1, 1)).
b_gmem_layout (Layout): Layout of the B global-memory tensor, used only by the CUDA-core fallback (defaults to Layout.row_major(1, 1)).

Args:

expert_usage_stats (Pointer[UInt32, ImmutAnyOrigin, _safe=False]): Pointer to per-expert usage stats; index 1 holds the active expert count.
a_tma_op (TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape]): TMA tile descriptor for loading A tiles from global to shared memory.
expert_ids (Pointer[Int32, ImmutAnyOrigin, _safe=False]): Pointer to the per-group expert id array; negative entries skip the tile.
b_tma_op (TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape]): TMA tile descriptor for loading B tiles from global to shared memory.
b_offsets (Pointer[UInt32, ImmutAnyOrigin, _safe=False]): Pointer to the per-expert B offset prefix sum, of length num_active_experts + 1.
c_tma_op (TMATensorTile[c_type, c_tile_rank, c_tile_shape_param, c_desc_shape]): TMA tile descriptor for storing C tiles from shared to global memory.
c_ptr (Pointer[Scalar[c_type], MutAnyOrigin, _safe=False]): Base pointer to the row-major C output tensor in global memory.
mnk (StaticTuple[UInt32, Int(3)]): [M, N, K] dimensions of the grouped GEMM as unsigned 32-bit integers.
a_gmem (LayoutTensor[a_type, a_gmem_layout, ImmutAnyOrigin]): A operand tensor used only by the CUDA-core fallback when use_tma is False.
b_gmem (LayoutTensor[b_type, b_gmem_layout, ImmutAnyOrigin]): B operand tensor used only by the CUDA-core fallback when use_tma is False.