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

batched_matmul_kernel_gpu

def batched_matmul_kernel_gpu[c_type: DType, a_type: DType, b_type: DType, CTensorType: TensorLayout, ATensorType: TensorLayout, BTensorType: TensorLayout, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b], elementwise_lambda_fn: Optional[def[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None] = None](c_tensor: TileTensor[c_type, CTensorType, MutAnyOrigin], a_tensor: TileTensor[a_type, ATensorType, ImmutAnyOrigin], b_tensor: TileTensor[b_type, BTensorType, ImmutAnyOrigin])

Computes a single batch slice of a batched matrix multiplication on the GPU by dispatching to the multistage GEMM kernel on NVIDIA or the AMDMatmul kernel on AMD hardware.

Parameters:

c_type (DType): Output tensor element dtype.
a_type (DType): LHS input tensor element dtype.
b_type (DType): RHS input tensor element dtype.
CTensorType (TensorLayout): Layout type of the output tensor.
ATensorType (TensorLayout): Layout type of the LHS input tensor.
BTensorType (TensorLayout): Layout type of the RHS input tensor.
transpose_b (Bool): Whether the RHS input is transposed.
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Matmul kernel configuration for the target hardware.
elementwise_lambda_fn (Optional[def[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None]): Optional epilogue applied to each output element.

Args:

c_tensor (TileTensor[c_type, CTensorType, MutAnyOrigin]): Rank-3 output tensor of shape (batch, m, n).
a_tensor (TileTensor[a_type, ATensorType, ImmutAnyOrigin]): Rank-3 LHS input tensor of shape (batch, m, k).
b_tensor (TileTensor[b_type, BTensorType, ImmutAnyOrigin]): Rank-3 RHS input tensor of shape (batch, k, n).