JAX Release 26.08

The NVIDIA JAX Release 26.08 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.4.1
  • Please refer to the CUDA DL 26.08 release notes section for the list of libraries inherited from the CUDA container.

• JAX jax-v0.11.0
• MaxText be69fa6
• Flax a31129e
• XLA 131bf41
• Transformer Engine 27486e03c

Driver Requirements

Release 26.08 is based on CUDA 13.4.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

MOE

• Integrated TransformerEngine Collective GEMM into MaxText for MoE layers, enabling GPU communication and computation overlap during training. (MaxText).
• Added fused MoE FFN1 combining wi_0 and wi_1 into a single grouped GEMM, reducing GPU kernel launch overhead for MoE MLP layers. (MaxText)
• Enabled overlapping MoE expert communications with collective matrix multiplication to improve GPU utilization during MoE forward and backward passes. (MaxText)
• Fixed ragged all-to-all buffer factor handling in DeepSeek-V3 MoE routing collectives. (MaxText)
• Added Expert Parallelism via a NCCL-EP backend in TransformerEngine for JAX, with bitmap top-k for faster MoE routing. (TransformerEngine)
• Extended communication+GEMM overlap API in JAX to support the cuBLASMp backend. (TransformerEngine)
• Made the linear offset of ClampedSwiGLU configurable in JAX for improved MoE activation control. (TransformerEngine)

Attention

• Added packed all-gather for context parallelism, improving GPU memory efficiency in long-context training. (MaxText)
• Added cuDNN-backed flex attention support in JAX. (TransformerEngine)

Compiler and Runtime

• JAX upgraded to v0.11.0 (from v0.10.2) with CUDA 13.4 support.
• Enabled experimental ragged all-to-all barrier with NCCL by default for improved GPU collective performance. (XLA)
• Added RAFT select_k for unstable TopK on GPU and lowered approx_top_k to stable sort+slice. (XLA)
• Integrated latest Triton updates and fixed thread pool starvation deadlock in GPU autotuning. (XLA)

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.