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 struct

BlackwellMatmulSM100FallbackKernel

struct BlackwellMatmulSM100FallbackKernel[a_type: DType, b_type: DType, c_type: DType, c_layout: TensorLayout, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], transpose_b: Bool = True, cluster_shape: StaticTuple[Int32, Int(3)] = StaticTuple(Int32(1), Int32(1), Int32(1)), a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, 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]

Simple fallback matmul kernel for SM100 (B200).

This kernel is used when the warp-specialized kernel is not applicable, such as for small problem sizes or unsupported configurations.

Unlike the main BlackwellMatmulSM100Kernel, this uses:

  • Single warp approach (no warp specialization)
  • Basic barrier synchronization (no CLC scheduling)
  • Direct TileTensor output (no TMA for C)
  • Simpler pipeline with single buffer

Parameters

  • a_type (DType): Element type of the A input matrix.
  • b_type (DType): Element type of the B input matrix.
  • c_type (DType): Element type of the C output matrix.
  • c_layout (TensorLayout): Memory layout of the C output tensor in global memory, used for output tiling and static stride computation.
  • block_tile_shape (IndexList[Int(3)]): Block tile dimensions (BM, BN, BK) for CTA-level tiling of the output and reduction dimensions.
  • mma_shape (IndexList[Int(3)]): MMA instruction dimensions (MMA_M, MMA_N, MMA_K) for the tensor core operation.
  • transpose_b (Bool): Whether B is stored transposed (K-major) (defaults to True).
  • cluster_shape (StaticTuple[Int32, Int(3)]): Thread block cluster dimensions used for LLVM cluster metadata (defaults to (1, 1, 1)).
  • a_swizzle (TensorMapSwizzle): Swizzle pattern for A shared memory tiles (defaults to TensorMapSwizzle.SWIZZLE_128B).
  • b_swizzle (TensorMapSwizzle): Swizzle pattern for B shared memory tiles (defaults to TensorMapSwizzle.SWIZZLE_128B).
  • num_threads (Int): Number of threads per CTA; must be 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 epilogue function applied to output elements (defaults to None).

Implemented traits

AnyType, Deinitable, Movable

comptime members

a_size

comptime a_size = (block_tile_shape[Int(0)] * block_tile_shape[Int(2)])

a_smem_layout_typed

comptime a_smem_layout_typed = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())))

a_swizzle_elems

comptime a_swizzle_elems = (a_swizzle.bytes() // size_of[a_type]())

accum_type

comptime accum_type = get_accum_type[a_type]()

ADescLayout

comptime ADescLayout = Layout[*?, *?]

ATile

comptime ATile = TileTensor[a_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]

ATileLayout

comptime ATileLayout = Layout[*?, *?]

ATmaOp

comptime ATmaOp = TMATensorTile[a_type, Int(2), _to_index_list[Layout[*?, *?]](), _to_index_list[Int(2), Layout[*?, *?]]()]

b_size

comptime b_size = (block_tile_shape[Int(1)] * block_tile_shape[Int(2)])

b_smem_layout_typed

comptime b_smem_layout_typed = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())))

b_swizzle_elems

comptime b_swizzle_elems = (b_swizzle.bytes() // size_of[b_type]())

BDescLayout

comptime BDescLayout = Layout[*?, *?]

BK

comptime BK = block_tile_shape[Int(2)]

BM

comptime BM = block_tile_shape[Int(0)]

BN

comptime BN = block_tile_shape[Int(1)]

BTile

comptime BTile = TileTensor[b_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]

BTileLayout

comptime BTileLayout = Layout[*?, *?]

BTmaOp

comptime BTmaOp = TMATensorTile[b_type, Int(2), _to_index_list[Layout[*?, *?]](), _to_index_list[Int(2), Layout[*?, *?]]()]

c_frag_size

comptime c_frag_size = (Int((mul mma_shape[Int(0)], mma_shape[Int(1)])) // num_threads)

CGmemStrideLayout

comptime CGmemStrideLayout = Layout[*?, *?]

max_tmem_cols

comptime max_tmem_cols = 512

MMA_K

comptime MMA_K = mma_shape[Int(2)]

MMA_M

comptime MMA_M = mma_shape[Int(0)]

MMA_N

comptime MMA_N = mma_shape[Int(1)]

num_k_mmas

comptime num_k_mmas = (block_tile_shape[Int(2)] // mma_shape[Int(2)])

num_m_mmas

comptime num_m_mmas = (block_tile_shape[Int(0)] // mma_shape[Int(0)])

num_n_mmas

comptime num_n_mmas = (block_tile_shape[Int(1)] // mma_shape[Int(1)])

static_N

comptime static_N = c_layout.static_stride[Int(0)]

Methods

validate_constraints

static def validate_constraints()

Validate compile-time constraints for this kernel configuration.

run

static def run(a_tma_op: TMATensorTile[a_type, Int(2), _to_index_list[Layout[*?, *?]](), _to_index_list[Int(2), Layout[*?, *?]]()], b_tma_op: TMATensorTile[b_type, Int(2), _to_index_list[Layout[*?, *?]](), _to_index_list[Int(2), Layout[*?, *?]]()], c: TileTensor[c_type, c_layout, MutAnyOrigin], num_iters: Int32)

Run the fallback matmul kernel.

Args: