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 package

nn

Provides neural network operators for deep learning models.

Packages

attention: Multi-head attention (MHA) and multi-head latent attention (MLA) kernels.
conv: Convolution kernels (1D, 2D, 3D, transposed).

Modules

activations: The module contains implementations of activation functions.
arange: Implements the arange operation that fills a tensor with evenly-spaced values over a range.
arg_nonzero: Implements arg_nonzero, which returns the indices of all nonzero elements in a tensor.
argmaxmin: Provides CPU and GPU implementations of the argmax and argmin reduction operations.
argmaxmin_gpu: Provides a GPU argmax/argmin kernel that wraps top-k with K=1 for inner-dimension reductions.
argsort: Provides CPU and GPU implementations of argsort that return indices permuting a tensor into sorted order.
bicubic: This module provides CPU and GPU implementations for bicubic interpolation.
broadcast: Implements numpy-style tensor broadcasting for CPU and GPU targets.
concat: Implements tensor concatenation along a specified axis for CPU and GPU targets.
cumsum: Implements the ONNX CumSum operator, computing prefix sums along a specified tensor axis.
fold: Implements the fold operation.
fused_qk_rope: Provides fused query/key rotary position embedding (RoPE) kernels with integrated KV-cache writes.
gather_scatter: Implements gather and scatter operations for CPU and GPU, including indexed reductions.
gemv_partial_norm: Fused matvec (M=1) + partial RMS norm on B200.
image: Provides pooling helpers and sliding-window shape utilities used by image-processing operations.
index_fp8: Implements tensor indexing and gather kernels with FP8 quantization support on Blackwell GPUs.
index_tensor: Implements numpy-style advanced tensor indexing (getitem and setitem) for CPU and GPU.
irfft: Inverse real FFT kernel using cuFFT.
kv_cache:
kv_cache_ragged: Implements KV-cache kernels for ragged (variable-length) sequences used in continuous batching.
learnable_2d_interp_pos_emb: Learnable 2D interpolated position embedding (Kimi K2.5 MoonViT3d).
moe: Implements Mixture-of-Experts (MoE) routing, token dispatch, and expert computation kernels.
nms: Implements non-maximum suppression (NMS) for object detection bounding-box filtering.
normalization:
pad: Implements tensor padding with constant or edge values for CPU and GPU.
pad_gpu: Implements GPU-specific tensor padding kernels with constant or edge-fill strategies.
pool: Provides average- and max-pooling kernels with configurable padding, dilation, and stride for CPU and GPU.
rand_normal: Generates tensors filled with values drawn from a normal (Gaussian) distribution for CPU and GPU.
rand_uniform: Generates tensors filled with values drawn from a uniform distribution for CPU and GPU.
randn: Provides a thin wrapper around randn to fill a TileTensor with standard-normal random values.
repeat_interleave: Implements repeat_interleave, which repeats each tensor element a specified number of times along an axis.
reshape: Implements the reshape operation that reinterprets a contiguous tensor's data under a new shape.
resize: Implements tensor resize (upsample/downsample) with nearest, bilinear, and other interpolation modes.
roi_align: Implements RoI Align, which extracts fixed-size feature maps from regions of interest using bilinear interpolation.
rope:
rope_split_store: Fused rope + split + KV store kernel.
sampling: Provides token-sampling utilities including logit-penalty application and repetition-penalty kernels.
shapes: Provides output-shape computation utilities for sliding-window operations such as pooling and convolution.
shard_and_stack: Implements shard-and-stack: shards a tensor across devices and stacks the shards into a higher-rank output.
slice: Implements the ONNX Slice operator, selecting sub-tensors along specified axes with start, stop, and step.
softmax: Provides numerically stable softmax kernels for CPU and GPU, including fused and online variants.
spatial_merge: Implements spatial merge, which compresses vision token grids by merging spatial blocks before attention.
split: Implements the tensor split operation, dividing a tensor into chunks along a specified axis.
tile: Implements the ONNX Tile operator, which replicates a tensor a specified number of times along each axis.
topk: Provides top-K selection kernels using warp- and block-level reductions for CPU and GPU.
topk_bitonic: Block-wide bitonic sort top-k for the MLA/MSA indexer.
topk_fi: Implements a float-integer packed top-K kernel that jointly tracks values and indices in a single register.
toppminp: Provides CPU implementations of top-p (nucleus) and min-p sampling for autoregressive token generation.
toppminp_gpu: Provides GPU implementations of top-p (nucleus) and min-p sampling for autoregressive token generation.
tpool_patch_merger: Implements the temporal-pool patch-merger kernel that reduces vision token sequences via average pooling.

Packages​

Modules​

Packages

Modules