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).

MAX accelerator library

The MAX accelerator library is the Mojo API for hardware-accelerated programming. It covers GPU primitives, the compute kernels that MAX graphs execute, the tensor layouts those kernels are written against, and the extension points for adding your own operations.

API summary

• Program a GPU directly with max.gpu: thread and block indexing, memory spaces, synchronization, and host-side device management.
• Describe how tensors sit in memory with layout: layouts, tiling, and the tensor types kernels are written against.
• Call ready-made compute kernels from linalg for linear algebra such as matrix multiplication, nn for neural-network operators such as attention and convolution, quantization for quantized weight encodings, and kv_cache for transformer key-value caches.
• Add your own operation to a MAX graph by writing it against extensibility.
• Spread work across GPUs and nodes with comm for multi-GPU collectives and shmem for multi-node communication.
• Time and profile a kernel with max.benchmark for timing harnesses and profiling_range for annotating profiler traces.

Package organization

max is a top-level package with subpackages, so its members carry the max prefix:

from max.gpu.compute import mma
from max.benchmark import Bench

The rest are top-level packages alongside max:

from linalg.matmul import matmul
from layout import TileTensor, row_major

The split tracks where the code lives today. More of these packages will move under max in later releases.

Packages

• ​algorithm: Free-form row-wise reduction library: the ReduceOp monoid trait, the rowwise scaffolder (Row / reduce / map / emit), and the Row-based pure reductions used by the graph reduction kernels.
• ​builtin_kernels: Provides the graph compiler's built-in kernel registrations.
• ​comm: Provides communication primitives for multi-GPU workloads.
• ​extensibility: Surface needed to write GraphCompiler kernel entry points.
• ​kv_cache: Contains implementations for several types of key-value caches.
• ​layout: Provides layout and layout tensor types, which abstract memory layout for multidimensional data.
• ​linalg: Provides CPU and GPU implementations of linear algebra functions.
• ​max: The MAX library for hardware-accelerated programming in Mojo.
• ​nn: Provides neural network operators for deep learning models.
• ​pipeline: Generic compile-time software pipelining for GPU kernels: schedule generation and verification.
• ​profiling_range: Provides Mojo wrappers for MAX's libkineto-backed range profiler.
• ​quantization: This package contains a set of APIs for quantizing tensor data.
• ​shmem: Implements a subset of OpenSHMEM for multi-node GPU communication.
• ​state_space: State space model kernels for Mamba and related architectures.
• ​structured_kernels: Shared GPU kernel primitives for structured kernel architectures.