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

test_matmul_sm90

def test_matmul_sm90[MType: CoordLike, NType: CoordLike, KType: CoordLike, //, a_type: DType, b_type: DType, c_type: DType, cluster_shape: IndexList[Int(3)], block_tile_shape: IndexList[Int(3)], wgmma_shape: IndexList[Int(3)], num_consumer: Int = Int(1), num_pipeline_stages: Int = Int(4), transpose_b: Bool = True, partitioned_multicast: Bool = False, grid_shape: OptionalReg[IndexList[Int(2)]] = None, use_tma_store: Bool = False, schedule: MatmulSchedule = MatmulSchedule.NONE, default_epilogue: Bool = False, elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]] = None, measure_threshold: Optional[Float64] = None, backend: Backend = Backend.CUBLAS, k_group_size: Int = Int(1)](ctx: DeviceContext, m: MType, n: NType, k: KType)

Runs an SM90 warp-specialized matmul and validates the result against a vendor reference.

Allocates host and device buffers for the operands, initializes them with random data, launches the warp-specialized GEMM with multicasting, and compares the output against a vendor (cuBLAS) reference using an elementwise tolerance check.

Parameters:

MType (CoordLike): Coordinate-like type encoding the M dimension extent (inferred).
NType (CoordLike): Coordinate-like type encoding the N dimension extent (inferred).
KType (CoordLike): Coordinate-like type encoding the K dimension extent (inferred).
a_type (DType): Element dtype of the left-hand operand.
b_type (DType): Element dtype of the right-hand operand.
c_type (DType): Element dtype of the output matrix.
cluster_shape (IndexList[Int(3)]): Thread block cluster shape as (N, M, K).
block_tile_shape (IndexList[Int(3)]): Block tile shape as (BM, BN, BK).
wgmma_shape (IndexList[Int(3)]): WGMMA instruction shape as (M, N, K).
num_consumer (Int): Number of consumer warps in the warp-specialized pipeline.
num_pipeline_stages (Int): Number of software pipeline stages.
transpose_b (Bool): Whether the right-hand operand is stored transposed.
partitioned_multicast (Bool): Whether to use partitioned instead of broadcast multicast.
grid_shape (OptionalReg[IndexList[Int(2)]]): Optional explicit grid shape override.
use_tma_store (Bool): Whether to use TMA store for the epilogue.
schedule (MatmulSchedule): Tile scheduler hint controlling launch order.
default_epilogue (Bool): Whether to use the default store-based epilogue.
elementwise_compute_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]]): Optional elementwise compute lambda applied in the epilogue.
measure_threshold (Optional[Float64]): Optional relative-difference threshold for a measured assertion.
backend (Backend): Vendor backend used for the reference matmul.
k_group_size (Int): Group size used for grouped matmul along the K dimension.

Args:

ctx (DeviceContext): Device context used for device allocations and kernel dispatch.
m (MType): Number of rows of the output matrix (M dimension).
n (NType): Number of columns of the output matrix (N dimension).
k (KType): Contraction dimension shared by the operands (K dimension).