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

warp_specialized_matmul

def warp_specialized_matmul[M: Int, N: Int, K: Int, BM: Int, BN: Int, BK: Int, WM: Int, WN: Int, WK: Int, a_producer_warps: Int, b_producer_warps: Int, consumer_warps: Int, pipeline_stages: Int = Int(1)](a_tt: TileTensor[DType.bfloat16, Storage=a_tt.Storage, address_space=a_tt.address_space, linear_idx_type=a_tt.linear_idx_type], b_tt: TileTensor[DType.bfloat16, Storage=b_tt.Storage, address_space=b_tt.address_space, linear_idx_type=b_tt.linear_idx_type], c_tt: TileTensor[DType.float32, Storage=c_tt.Storage, address_space=c_tt.address_space, linear_idx_type=c_tt.linear_idx_type], ctx: DeviceContext)

Enqueues the warp-specialized matrix multiplication kernel onto the device.

Converts the input and output TileTensor arguments to global address-space layout tensors and launches the kernel with a grid of M // BM by N // BN blocks and one warp per producer and consumer warp.

Parameters:

M (Int): Number of rows of the output matrix, sets the grid height M // BM.
N (Int): Number of columns of the output matrix, sets the grid width N // BN.
K (Int): Contraction dimension of the matmul.
BM (Int): Block tile size along the M dimension.
BN (Int): Block tile size along the N dimension.
BK (Int): Block tile size along the K dimension.
WM (Int): Warp tile size along the M dimension.
WN (Int): Warp tile size along the N dimension.
WK (Int): Warp tile size along the K dimension.
a_producer_warps (Int): Number of warps assigned to loading matrix A tiles.
b_producer_warps (Int): Number of warps assigned to loading matrix B tiles.
consumer_warps (Int): Number of warps assigned to computing the matmul.
pipeline_stages (Int): Number of ring buffer stages used to overlap loads and compute (defaults to 1).

Args:

a_tt (TileTensor[DType.bfloat16, Storage=a_tt.Storage, address_space=a_tt.address_space, linear_idx_type=a_tt.linear_idx_type]): Input matrix A as a TileTensor of shape M x K.
b_tt (TileTensor[DType.bfloat16, Storage=b_tt.Storage, address_space=b_tt.address_space, linear_idx_type=b_tt.linear_idx_type]): Input matrix B as a TileTensor of shape K x N.
c_tt (TileTensor[DType.float32, Storage=c_tt.Storage, address_space=c_tt.address_space, linear_idx_type=c_tt.linear_idx_type]): Output matrix C as a TileTensor of shape M x N.
ctx (DeviceContext): Device context used to enqueue the kernel.