JAX Release 26.05
The NVIDIA JAX Release 26.05 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.2.1
- Please refer to the CUDA DL 26.05 release notes section for the list of libraries inherited from the CUDA container.
Driver Requirements
Release 26.05 is based on CUDA 13.2.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
- jax.ragged_dot has better performance. It's powered by XLA-cuDNN. Use sparse_matmul=true in MaxText to invoke jax.ragged_dot.
- cuBLASLt is an alternative gemm backend, replacing cuBLAS. The autotuning of gemms now spans across Triton, cuDNN and cuBLASLt.
- Enabled deviceless XLA-cuDNN fusion for AoT workflow to reduce compilation time.
- Expanded fusion support of convolution with XLA-cuDNN graph API integration with NHWC layout.
- Added CUPTI 13.2 multi-subscriber tracing, allowing XLA profiling to run in the same process as other CUPTI-based tools such as NSight.
Performance Optimizations
- Accelerated jax.ragged_dot with XLA-cuDNN integration. fprop/dgrad powered by ragged gemm while wgrad is by batched dense gemm. Use xla_gpu_experimental_use_ragged_dot_fusion = ON.
- Added basic PDL (Programmatic Dependent Launch) in XLA, uplifting E2E performance of inference models.For example, Gemma3’s perf uplifted by single digit percentage.
- Introduced native kernel codegen improvements for FP4 training, including all-gather-transpose optimization, higher unroll/vectorization factors for narrow data types, and lowering slice/slice-update from FP4 activation offloading to D2H memcpy.
- Enabled async multi-stream collectives in XLA, permitting inter-node communication overlap with that of intra-node communication. For example, Llama3-405B perf uplifted by single digit percentage.
- Enhanced XLA to leverage NCCL 2.28’s copy engine to reduce SM usage of XLA’s all-gather & all-2-all communication collectives. For example, MaxText’s Llama3 perf uplifted by single digit percentage.
- Improved D2H & H2D copy overlap with compute, to help accelerate MoE models on GB200/GB300 systems.
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.
| Container | Type | Image URI |
|---|---|---|
| jax | - | ghcr.io/nvidia/jax:jax-YYYY-MM-DD |
| maxtext | LLM framework | ghcr.io/nvidia/jax:maxtext-YYYY-MM-DD |
| axlearn | LLM framework | ghcr.io/nvidia/jax:axlearn-YYYY-MM-DD |
Known Issues
- Transformer Engine version 2.14 shipped with 26.04 has an issue resulting in a nondeterministic wrong answer when using MXFP8 training and bias is present. Please upgrade to Transformer Engine version 2.14.1, where the issue was fixed.