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 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:

Models

Cookbook

Puzzles

Blog

Community

Contact