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

select_config

def select_config[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = False](M: Int, N: Int, K: Int, ctx: DeviceContext) -> MatmulConfig[a_type, b_type, c_type, transpose_b]

Selects a heuristic-optimal MatmulConfig for the given problem shape and device.

Evaluates candidate block tile shapes and split-K partition counts, then chooses the configuration that minimizes estimated work per streaming multiprocessor while keeping the wave count bounded.

Parameters:

a_type (DType): The DType of the left-hand operand A.
b_type (DType): The DType of the right-hand operand B.
c_type (DType): The DType of the output C.
transpose_b (Bool): Whether B is supplied transposed (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.
ctx (DeviceContext): The device context used to query GPU properties.

Returns:

MatmulConfig[a_type, b_type, c_type, transpose_b]: The selected MatmulConfig for the given problem.