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

choose_config

def choose_config[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = True, gemm_kind: GEMMKind = GEMMKind.GEMM, has_epilogue_tensor: Bool = False, epilogue_is_1d: Bool = False](M: Int, N: Int, K: Int, B: Int) -> MatmulConfig[a_type, b_type, c_type, transpose_b]

Select a MatmulConfig that minimizes waves per SM for the given problem shape.

Parameters:

a_type (DType): DType of the A (left) operand elements; must equal b_type.
b_type (DType): DType of the B (right) operand elements.
c_type (DType): DType of the output matrix elements.
transpose_b (Bool): Whether the B operand is stored transposed (defaults to True).
gemm_kind (GEMMKind): The GEMMKind selecting the kernel variant, for example GEMM or BMM (defaults to GEMMKind.GEMM).
has_epilogue_tensor (Bool): Whether the kernel uses a TMA epilogue load for an epilogue tensor (defaults to False).
epilogue_is_1d (Bool): Whether the epilogue tensor is 1D, for example a bias vector (defaults to False).

Args:

M (Int): The M dimension of the matmul.
N (Int): The N dimension of the matmul.
K (Int): The K dimension of the matmul.
B (Int): The batch dimension of the matmul.

Returns:

MatmulConfig[a_type, b_type, c_type, transpose_b]: A MatmulConfig tuned for the given problem dimensions.