JAX Release 26.07

The NVIDIA JAX Release 26.07 is made up of two container images available on NGC: JAX and MaxText.

Contents of the JAX container

This container image contains the complete source for the following software:

• JAX: /opt/jax
• XLA: /opt/xla
• Flax:/opt/flax
• TransformerEngine:/opt/transformer-engine

The MaxText container image is based on the JAX container. Additionally, it includes:

• MaxText: /opt/maxtext

The JAX runtime package jaxlib is prebuilt and installed in the default Python environment (/usr/local/lib/python3.10/dist-packages/jaxlib) in the container image.

Versions of packages included in both of these containers:

• CUDA 13.3 U1
  • Please refer to the CUDA DL 26.07 release notes section for the list of libraries inherited from the CUDA container.

• JAX jax-v0.10.2
• MaxText 454e55f
• Flax 5768bb5
• XLA 5a9e73c (internal)
• Transformer Engine 4220403
• AXLearn c8624c8

Driver Requirements

Release 26.07 is based on CUDA 13.3.1. For comprehensive and up-to-date driver compatibility information, please refer to the following documentation:

• NVIDIA CUDA Compatibility Guide - Compatibility information between CUDA versions and driver releases.
• CUDA Toolkit Release Notes - Driver version requirements and compatibility matrices.
• NVIDIA Drivers Download - Latest NVIDIA drivers.

Key Features and Enhancements

• The GCP NCCL plugin package is now installed in the container. On Google Cloud Platform hosts (A3 Ultra, A4, A4X and A4X Max machine types ), the container automatically detects and activates the GCP NCCL network, tuner, profiler, and environment plugins for optimized multi-node communication.
  • On GCP hosts with the GCP NCCL plugins installed, NCCL telemetry is enabled by default starting in this release.
  • A runtime advisory message is displayed when running on GCP hosts with the GCP NCCL plugins active, directing users to the GCP NCCL telemetry documentation.
• PDL launch inserted after Triton GEMM MMA loops before epilogues by default, increasing producer-consumer overlap with e2e speed up on smaller models.
• cuDNN-backed flex attention in jax.dot_product_attention reenabled.
• Support for CuDNN’s AoT/deviceless compilation in XLA. cuDNN backend invoked in deviceless mode in XLA, improving compilation time for graphs with cuDNN calls.
• Fixed incorrect gradients in zero-padded positions in jax.dot_product_attention.
• Redundant zero-initialization removal in DS/DUS scan-loops, giving small end to end speedups in LLM model training.

Deprecations

• Removal of NVSHMEM calls/library dependency from XLA. Developers using NVSHMEM in custom kernels are recommended to leverage NCCL device-side APIs.

JAX Toolbox

The JAX Toolbox projects focus on achieving the best performance and convergence on NVIDIA Ampere, Hopper, and Blackwell architecture families and provide the latest deep learning models and scripts for training and fine-tuning. These examples are tested against a nightly CI as well as each NGC container release to ensure consistent accuracy and performance over time.

Nightly Containers

In addition to projects, JAX Toolbox includes nightly containers for libraries across the JAX ecosystem.

Known Issues

• None.