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

multistage_gemm

def multistage_gemm[c_type: DType, a_type: DType, b_type: DType, //, *, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b], elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[b_type, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type], ctx: DeviceContext)

TileTensor overload of multistage_gemm. Converts to LayoutTensor and dispatches to a GEMM kernel.

Parameters:

c_type (DType): DType of the output tile c elements (inferred).
a_type (DType): DType of the input tile a elements (inferred).
b_type (DType): DType of the input tile b elements (inferred).
transpose_b (Bool): Whether b is accessed transposed, so its row y is read as b[y, i] instead of b[i, y].
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Compile-time MatmulConfig selecting the block tile, warp tile, MMA shape, and pipeline stages for the kernel.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue applied to the accumulated value before it is stored to c (defaults to None, which stores the raw accumulation).

Args:

c (TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output tile of shape (M, N) receiving the matmul result.
a (TileTensor[a_type, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input tile of shape (M, K).
b (TileTensor[b_type, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Input tile of shape (K, N), or (N, K) when transpose_b is set.
ctx (DeviceContext): Device context used to enqueue the kernel.

TileTensor overload of multistage_gemm with runtime config. Constrains c to mut=True because split_k_reduce requires a mutable output tensor.

Parameters:

c_type (DType): DType of the output tile c elements (inferred).
a_type (DType): DType of the input tile a elements (inferred).
b_type (DType): DType of the input tile b elements (inferred).
transpose_b (Bool): Whether b is accessed transposed, so its row y is read as b[y, i] instead of b[i, y].
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Compile-time MatmulConfig selecting the block tile, warp tile, MMA shape, and pipeline stages for the kernel.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue applied to the accumulated value before it is stored to c (defaults to None, which stores the raw accumulation).

Args:

c (TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output tile of shape (M, N) receiving the matmul result; must be mutable because split_k_reduce writes the reduced sum back into it.
a (TileTensor[a_type, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input tile of shape (M, K).
b (TileTensor[b_type, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Input tile of shape (K, N), or (N, K) when transpose_b is set.
runtime_config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Runtime MatmulConfig carrying the number of K partitions used for the split-K reduction path; when num_k_partitions is greater than 1 the kernel writes partial sums to a workspace and reduces them into c.
ctx (DeviceContext): Device context used to enqueue the kernel and allocate the split-K workspace.