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

def get_partitioned_matmul_mojo[b_type: DType, kernel_rows: Int, kernel_cols: Int, use_i8mm: Bool = False](m: Int, n: Int, k: Int, task_id: Int, num_tasks: Int) -> SubMatmulConfig

Returns the sub-matmul config for a task using the Mojo partitioner.

Splits the work into row and column tasks via partition_work after determining the task grid with get_partitioned_matmul_mojo_shape.

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.
task_id (Int): The index of the task to compute the partition for.
num_tasks (Int): The total number of tasks sharing the matmul.

Returns:

SubMatmulConfig: The sub-matmul offset and shape for this task.