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_kernel

def warp_specialized_matmul_kernel[in_type: DType, out_type: DType, a_layout: Layout, b_layout: Layout, c_layout: Layout, 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](a: LayoutTensor[in_type, a_layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL], b: LayoutTensor[in_type, b_layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL], c: LayoutTensor[out_type, c_layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL])

Runs the warp-specialized matrix multiplication kernel on an AMD GPU.

Splits the block's warps into A producers, B producers, and consumers, coordinating tile movement through ring buffers with pipeline stages. Producers load BM x BK and BN x BK tiles from global memory into shared memory, while consumers accumulate the matrix multiply-accumulate results and write them back to global memory.

Parameters:

in_type (DType): Element type of the input matrices a and b.
out_type (DType): Element type of the output matrix c.
a_layout (Layout): Memory layout of input matrix a in global memory.
b_layout (Layout): Memory layout of input matrix b in global memory.
c_layout (Layout): Memory layout of output matrix c in global memory.
BM (Int): Block tile size along the M dimension, must be divisible by WM.
BN (Int): Block tile size along the N dimension, must be divisible by WN.
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.

Args:

a (LayoutTensor[in_type, a_layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL]): Input matrix A as a global memory tensor of shape M x K.
b (LayoutTensor[in_type, b_layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL]): Input matrix B as a global memory tensor of shape K x N.
c (LayoutTensor[out_type, c_layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL]): Output matrix C as a global memory tensor of shape M x N.