For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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

ElementwiseComputeFn

comptime ElementwiseComputeFn = def[dtype: DType, width: SIMDLength, *, alignment: Int](IndexList[Int(2)], SIMD[dtype, width]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable

Value-taking counterpart of elementwise_compute_lambda_type.

A unified closure passed as a runtime value carries the origins of its captures, so the buffers it reads stay alive until the kernel that calls it has launched. alignment has no default; pass alignment=1 where the legacy type relied on its default.

ElementwiseEpilogueFn

comptime ElementwiseEpilogueFn = def[dtype: DType, width: SIMDLength, *, alignment: Int](IndexList[Int(2)], SIMD[dtype, width]) -> None & RegisterPassable & ImplicitlyCopyable

Value-taking counterpart of elementwise_epilogue_type.

The closure stores the output itself, so a kernel that receives one writes nothing to its output argument. Pass it to a named kernel with host_arg=. alignment has no default; pass alignment=1 where the legacy type relied on its default.

ElementwiseOutputComputeFn

comptime ElementwiseOutputComputeFn = def[dtype: DType, width: SIMDLength, *, alignment: Int](IndexList[Int(2)], SIMD[dtype, width], SIMD[dtype, width]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable

An ElementwiseComputeFn that also receives the output's prior value.

Called as fn(idx, val, c_val), where c_val is the output tensor's value at idx before the kernel writes it, cast to dtype. The kernel reads c_val from its own output argument, so the closure never captures the output and cannot alias it.

no_epilogue_fn

comptime no_epilogue_fn = _as_epilogue_fn[inflated$def[_no_epilogue_body[_, _, _]]](.__init__())

Stands in for an ElementwiseEpilogueFn argument on code paths that store the output directly.

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_elementwise_epilogue: Applies the legacy epilogue lambda if set, otherwise epilogue_fn.
• ​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.
• ​identity_compute_fn: Returns val unchanged.
• ​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.