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

get_partitioned_matmul_mojo_shape

def get_partitioned_matmul_mojo_shape[b_type: DType, kernel_rows: Int, kernel_cols: Int, use_i8mm: Bool](m: Int, n: Int, k: Int, num_tasks: Int) -> IndexList[Int(2)]

Returns the row and column task counts that best balance a parallel matmul.

Searches over feasible task grids and selects the one that minimizes the per-task work, with heuristics for small m and L2-cache-aware column partitioning.

Parameters:

b_type (DType): The dtype of the B tensor.
kernel_rows (Int): The static number of rows of the micro kernel.
kernel_cols (Int): The static number of columns of the micro kernel.
use_i8mm (Bool): Whether the i8mm instruction is available on the target.

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.
num_tasks (Int): The total number of tasks sharing the matmul.

Returns:

IndexList[Int(2)]: A pair (num_row_tasks, num_col_tasks) describing the task grid.