Transformer Engine v2.17 Release Notes

Key Features and Enhancements

  • [Common, JAX, PyTorch] Added Expert Parallelism support using a NCCL-EP backend across C API, PyTorch and JAX. (#3127) (#3035) (#3036)

  • [Common, JAX, PyTorch] Added a bitmap top-k implementation for faster MoE routing when using TE Expert Parallelism. (#3009)

  • [Common] Optimized the fused router forward/backward kernels, particularly for large expert counts. (#3012), (#2776)

  • [Common, PyTorch] Reduced the CPU overhead of the fused MLP block (#3001)

  • [Common, JAX, PyTorch] Extended the communication+GEMM overlap API to support the cuBLASMp backend. (#2443)

  • [Common] Added columnwise-only output support for NVFP4 2D block scaling.. (#3027)

  • [Common, PyTorch, JAX] Made the linear offset of ClampedSwiGLU configurable. (#2938)

  • [Common, PyTorch] Added the 4over6 option to NVFP4 recipe. (#2972)

  • [PyTorch] Enabled NVFP4 fused grouped MLP. (#3048)

  • [PyTorch] Added public utilities to interleave and de-interleave fused grouped-MLP FC1 weights for checkpoint conversion. (#3078)

  • [PyTorch] Added CPU-offloading support for grouped linear and fused grouped MLP operations. (#3047)

  • [PyTorch] Made GroupedLinear module compatible with CUDA graph capture. (#3038)

  • [PyTorch, JAX] Added support for cuDNN-backed flex attention. (#2984, #2985)

  • [PyTorch] Added FlashAttention support for MLA in THD format by padding unequal Q/V head dimensions.(#2629)

  • [JAX] Added a Flax MoEBlock composing fused routing, token permutation, grouped GEMMs, and communication. (#2912)

  • [Docs] Added a GEMM Profiling Guide and companion benchmark tool for BF16, FP8 block scaling, MXFP8, and NVFP4. (#2863)

Fixed Issues

  • [Common] Fixed an int32 overflow in multi_tensor_apply tensor sizes for numel > INT_MAX. (#3136)

  • [Common] Fixed a CUTLASS grouped GEMM crash when every group was empty. (#3067)

  • [Common] Guarded nvmlGpuFabricInfo_v2 (MNNVL) usage to avoid failures on unsupported systems. (#3013)

  • [Common] Prevented fused attention from being selected for return_max_logit=True with cuDNN versions earlier than 9.21, where that output combination is unsupported. (#3091)

  • [Common, PyTorch] Disabled cuDNN 9.23.0/9.23.1 for MXFP8 attention to avoid known issues. (#3173)

  • [PyTorch] Enabled FlashAttention 2 for head dimensions up to 256 on sm103 and other architectures. (#2836)

  • [PyTorch] Fixed wrong stream capture for the wgrad GEMM in the fused MoE MLP, which could corrupt gradients or produce NaNs under FSDP. (#3089)

  • [PyTorch] Removed redundant nested loops in GroupedLinear FP8 calibration so each input and weight is calibrated once per GEMM. (#3101)

  • [PyTorch] Fixed the scale_inv_min debug statistic incorrectly reporting zero for padded MXFP8 and NVFP4 scale tensors. (#3041)

  • [JAX] Fixed JAX logical sharding constraints under shard_map when Flax logical-axis rules are active. (#3103)

Breaking Changes in This Release

There are no breaking changes in this release.

Deprecated Features

There are no deprecated features in this release.