IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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], a_tma_op: TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape], expert_ids: Pointer[Int32, ImmutAnyOrigin], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape], b_offsets: Pointer[UInt32, ImmutAnyOrigin], c_tma_op: TMATensorTile[c_type, c_tile_rank, c_tile_shape_param, c_desc_shape], c_ptr: Pointer[Scalar[c_type], MutAnyOrigin], 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:

Models

Cookbook

Puzzles

Blog

Community

Contact