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

naive_batched_matmul_kernel

def naive_batched_matmul_kernel[rank: Int, c_type: DType, a_type: DType, b_type: DType, CTensorType: TensorLayout, ATensorType: TensorLayout, BTensorType: TensorLayout, transpose_b: Bool = False, 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, accum_type: DType = get_accum_type[c_type]()](c_tensor: TileTensor[c_type, CTensorType, MutAnyOrigin], a_tensor: TileTensor[a_type, ATensorType, ImmutAnyOrigin], b_tensor: TileTensor[b_type, BTensorType, ImmutAnyOrigin], c_buff_nd_shape: IndexList[rank])

Computes one element per thread of a batched matrix multiplication using a naive scalar accumulation loop over the contraction dimension.

Parameters:

rank (Int): Rank of the original (un-collapsed) output tensor.
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.
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.
accum_type (DType): Accumulator dtype used during the contraction.

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).
c_buff_nd_shape (IndexList[rank]): Shape of the original output tensor before collapsing to 3D, used to un-collapse batch coordinates for the epilogue.