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

Python module

max.nn

Neural network modules for MAX.

max.nn layers come in two modes. In graph-based execution, a layer’s __call__ stages its operations into a Graph, and MAX compiles the finished graph into an optimized executable that runs on Buffer inputs:

import numpy as np

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn import Linear

input_type = TensorType(DType.float32, ["batch", 64], device=DeviceRef.CPU())
with Graph("layer", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor
    graph.output(Linear(64, 32, DType.float32, DeviceRef.CPU())(x))

model = InferenceSession().load(
    graph,
    weights_registry={"weight": np.zeros((32, 64), dtype=np.float32)},
)

In eager execution, max.experimental.nn runs the same layer abstractions directly on Tensor values, executing each operation immediately, PyTorch-style.

Submodules

max.nn.attention	The attention mechanism used within the model.
max.nn.comm
max.nn.kernels	Helper functions for wrapping custom kv cache/attention related ops.
max.nn.kv_cache	Key-value cache types for MAX layers.
max.nn.moe	Mixture of Experts (MoE) module.
max.nn.sampling	Sampling custom ops.
max.nn.state_space	Python wrappers for the state_space Mojo kernel package.

Base classes

Identity	Identity layer that passes through input unchanged.
Layer	Base class for neural network components.
LayerList	Stores a list of layers.
MissingLayer	Holds the place of a decoder layer another pipeline stage builds.
Module	Base class for model components with weight management.
Sequential	A sequential stack of layers.
Shardable	Protocol for objects that support sharding across multiple devices.
Signals	Signal buffers used for peer-to-peer communication in allreduce.

Linear layers

ColumnParallelLinear	A Linear layer where the weight and bias are sharded onto multiple devices.
Embedding	A lookup table for embedding integer indices into dense vectors.
FusedMLP	Stores the gate and up projections as one pre-fused weight.
GPTQLinear	A Linear layer for GPTQ encoding.
Linear	Applies a linear transformation to incoming data: $y = xW^T + b$.
MLP	Simple multi-layer perceptron composed of three Linear layers.
StackedLinear	A module that manages multiple linear projections as a stacked weight.
VocabParallelEmbedding	A lookup table for embedding integer indices into dense vectors.

Normalization

ConstantLayerNorm	Layer normalization block with constant gamma and beta values.
GroupNorm	Group normalization block.
LayerNorm	Layer normalization over the last dimension.
RMSNorm	Computes the root mean square normalization on inputs.

Rotary embeddings

DynamicRotaryEmbedding	Applies RoPE with dynamic scaling for long-context inference.
LinearScalingParams	Scaling parameters for linear RoPE frequency scaling.
Llama3RopeScalingParams	Scaling parameters for Llama3's frequency-based context extension.
Llama3RotaryEmbedding	Applies RoPE with Llama3-style frequency scaling for extended context lengths.
LongRoPERotaryEmbedding	Applies RoPE with LongRoPE scaling for Phi-3.5 models.
LongRoPEScalingParams	Parameters for LongRoPE scaling as used in Phi-3.5 models.
RotaryEmbedding	Applies Rotary Position Embedding (RoPE) to transformer activations.
YarnRotaryEmbedding	Applies generic YaRN (Yet another RoPE eNhancement) Rotary Position Embedding.
YarnScalingParams	Scaling parameters for YaRN (Yet another RoPE eNhancement) frequency interpolation.

Transformer

DistributedTransformer	Transformer model consisting for TransformerBlock layers.
DistributedTransformerBlock	Stack of Attention, FeedForward, and RMSNorm layers.
HyperConnection	Mixes hc_mult parallel residual streams around a transformer sublayer.
ReturnHiddenStates	Which hidden states a transformer model returns per sequence.
ReturnLogits	Which logits a transformer model returns per sequence.
Transformer	A transformer model consisting of TransformerBlock layers.
TransformerBlock	Stack of Attention, FeedForward, and RMSNorm layers.
DistributedLogitsPostprocessMixin	Mixin providing logits postprocessing for multi-device sharded models.
LogitsPostprocessMixin	Mixin providing logits postprocessing for single-device models.

distributed_logits_postprocess	Common logits postprocessing for multi-device sharded models.
forward_sequential_layers	Forward pass through sequential layers with optional subgraph groups.
logits_postprocess	Performs common logits postprocessing for single-device models.

Convolution

Conv1D	A 1D convolution over an input signal composed of several input planes.
Conv2d	A 2D convolution over an input signal composed of several input planes.
Conv3D	A 3D convolution over an input signal composed of several input planes.
ConvTranspose1d	A 1D transposed convolution operator over an input image composed of several input planes.
WeightNormConvTranspose1d	A 1D transposed convolution operator over an input image composed of several input planes.

Mixture of experts

Allreduce	Layer to perform allreduce operation with automatic implementation selection.
EPCommBuffers	SHMEM communication buffers for one Expert Parallelism MoE forward pass.
MoE	Implementation of Mixture of Experts (MoE).
MoEGate	Gate module for MoE.
MoEQuantized	Mixture of Experts with FP8 or NVFP4 quantization.

Parallelism

ParallelArrayOps	Parallelized numpy array operations for performance-critical data processing.

Sampling

MinPSampler	A min_p sampler.
RejectionSampler	Rejection sampler for speculative decoding verification.
RejectionSamplerWithResiduals	A simple rejection sampler.

Quantization

QuantConfig	Configures scaled quantization settings for a layer or model section.
QuantFormat	Identifies the quantization format of a model checkpoint.
InputScaleSpec	Specifies how input activations are scaled for scaled quantization.
ScaleGranularity	Specifies the granularity of the quantization scale factor.
ScaleOrigin	Specifies whether the quantization scale is determined statically or dynamically.
WeightScaleSpec	Specifies how weights are scaled for scaled quantization.

Hooks

PrintHook	Hook that prints layer tensor inputs and outputs.

add_layer_hook	Adds a hook to call a function after each layer's __call__.
clear_hooks	Remove all hooks.

Functions

build_max_lengths_tensors	Builds two [1] uint32 scalar buffers of maximum lengths.
clamp	Clamps values in x to [min, max].
forward_moe_sharded_layers	Forward pass through DP-sharded layers (EP MoE or replicated MLP/MoE).
split_batch	Split a ragged input batch into data parallel batches.
split_batch_replicated	Split a ragged token batch into data parallel batches.