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:
arg_nonzero:
argmaxmin:
argmaxmin_gpu:
argsort:
bicubic: This module provides CPU and GPU implementations for bicubic interpolation.
broadcast:
concat:
cumsum:
fold: Implements the fold operation.
fused_qk_rope:
gather_scatter:
gemv_partial_norm: Fused matvec (M=1) + partial RMS norm on B200.
image:
index_fp8:
index_tensor:
irfft: Inverse real FFT kernel using cuFFT.
kv_cache:
kv_cache_ragged:
learnable_2d_interp_pos_emb: Learnable 2D interpolated position embedding (Kimi K2.5 MoonViT3d).
moe:
nms:
normalization:
pad:
pad_gpu:
pool:
rand_normal:
rand_uniform:
randn:
repeat_interleave:
reshape:
resize:
roi_align:
rope:
rope_split_store: Fused rope + split + KV store kernel.
sampling:
shapes:
shard_and_stack:
slice:
softmax:
spatial_merge:
split:
tile:
topk:
topk_fi:
toppminp:
toppminp_gpu:
tpool_patch_merger:

Packages​

Modules​

Packages

Modules