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 module

utils

Provides shared CPU matmul utilities including tile-consumer traits, kernel shape selection, and partial SIMD load/store helpers.

`comptime` values

`elementwise_compute_lambda_type`

comptime elementwise_compute_lambda_type = def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]

`elementwise_epilogue_type`

comptime elementwise_epilogue_type = def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None

Structs

GemmShape: Helper class to unpack gemm dimension and layout.
InnerKernelID: Identifies the inner matmul kernel variant selected for a target.
KernelConfig: Static configuration of the matmul inner kernel.
MicroKernelShape: Record describing the inner kernel shape.
NullTileConsumer: No-op TileConsumer. Used as the default when no fusion is requested, and as the placeholder type when a kernel's tile_consumer Optional is None.
NullTileOperation: No-op TileOperation sentinel: parallel to NullTileConsumer. Used as the default TileOperationType so kernels without a fused op compile without callers having to spell out a placeholder.
SubMatmulConfig: Static configuration of sub-matrices in parallel matmul.

Traits

TileConsumer: Trait for an epilogue operation which consumes a tile of data.
TileOperation: Non-terminal counterpart to TileConsumer: mutates a tile in place between MMA and store. Composes additively with capability subtraits like AuxLoading for ops that need pre-loaded data.

Functions

apply_epilogue: Applies an elementwise epilogue lambda to a source tile.
calculate_tile_n_k: Helper heuristic function to decide on tile size to partition the matmul given the cache size and desired data layout.
dispatch_get_kernel_type: Invokes func with the kernel type selected for the matmul dimensions.
get_kernel_config: Extracts matmul configuration parameters for exported functions.
get_kernel_type: Returns the kernel shape variant to use based on the matmul dimensions.
get_matmul_arch_factor: Returns the architecture-dependent alignment factor for matmul tiling.
get_matmul_kernel_shape: Returns the micro kernel shape for the current target and dtypes.
get_matmul_kernel_shape_ARM: Returns the micro kernel shape tuned for ARM targets.
get_matmul_kernel_shape_x86: Returns the micro kernel shape tuned for x86 targets.
get_matmul_num_tasks: Compute the number of tasks for parallel matmul. The max number of tasks is the thread or core count.
get_matmul_prefetch_b_distance_k: Returns the K-dimension prefetch distance for the B matrix.
get_min_task_size: Returns the minimum task size used to limit parallel matmul task counts.
get_pack_data_size: Utility to compute the number of elements to pack in each tile.
get_packB_unroll_factor: Returns the unroll factor applied while packing the B matrix.
get_partitioned_matmul: Returns the sub-matmul config for a given task in a parallel matmul.
get_partitioned_matmul_mojo: Returns the sub-matmul config for a task using the Mojo partitioner.
get_partitioned_matmul_mojo_shape: Returns the row and column task counts that best balance a parallel matmul.
is_valid_epilogue: Whether T is a real epilogue rather than a null sentinel.
lora_qkv_plane_row_offset: Additive row offset into a 3-plane row-stacked activation [3*M, K].
packA_i8mm: Packs a range of rows of matrix A for the i8mm kernel layout.
partial_simd_load: Loads a vector with dynamic bound.
partial_simd_store: Stores a vector with dynamic bound.
partition_work: Partitions work into blocks distributed across num_tasks tasks.
select_inner_kernel: Returns the inner kernel variant to use for the given dtypes and target.
use_i8mm_fn: Returns whether the i8mm instruction should be used for the given dtypes.
use_vnni_fn: Returns whether the VNNI instruction should be used for the given dtypes.

comptime values​

elementwise_compute_lambda_type​

elementwise_epilogue_type​

Structs​

Traits​

Functions​