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 toTrue). - 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 toTensorMapSwizzle.SWIZZLE_128B). - b_swizzle (
TensorMapSwizzle): Swizzle pattern for B shared memory tiles (defaults toTensorMapSwizzle.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 toNone).
Implemented traits
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:
- a_tma_op (
TMATensorTile[a_type, Int(2), _to_index_list[Layout[*(), *()]](), _to_index_list[Int(2), Layout[*(), *()]]()]): TMA descriptor for matrix A. - b_tma_op (
TMATensorTile[b_type, Int(2), _to_index_list[Layout[*(), *()]](), _to_index_list[Int(2), Layout[*(), *()]]()]): TMA descriptor for matrix B. - c (
TileTensor[c_type, c_layout, MutAnyOrigin]): Output tensor C (TileTensor, direct global memory writes). - num_iters (
Int32): Number of K-dimension iterations.