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 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 single-pass GPU argmax/argmin kernel over the inner-most dimension.
• ​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.