Transformer Engine v2.21 Release Notes

Key Features and Enhancements

• [PyTorch] Added experimental cuDNN-backed KimiDeltaAttention. (#3548)
• [Common, PyTorch] Added prototype NVFP4 quantization with UE5M3 block scales. (#3325)
• [PyTorch] Expanded torch.compile(fullgraph=True) support to FP16 and BF16 FusedAttention, fusible BasicLinear and Bias operations including Linear + Bias forward fusion, and FP8 weight caching in compiled Linear calls using is_first_microbatch. (#3472) (#3547) (#3505)
• [PyTorch] Added FlashAttention 3 and FlashAttention 4 support for THD attention with pad_between_seqs=True, including P2P and all-to-all context parallelism. (#3205)
• [PyTorch] Added support for FlashAttention 2.8.4. (#3403)
• [PyTorch] Added support for backward with head dimension 256 in THD and context parallel attention when using FlashAttention 4.0.0b33+ (#3540)
• [PyTorch] Improved THD training performance with dropout on SM100 and SM103. (#3313)
• [PyTorch] Added MoeDispatch and MoeCombine fusible operations for building MoE blocks with te.ops.Sequential and improved the performance of it. (#3503)
• [Common, PyTorch] Added a deterministic mode to fused_moe_aux_loss through deterministic=True. (#3587)
• [PyTorch] Extended the fused grouped MLP with deterministic dprob for GLU activations and fused activation recompute for ScaledTanhSReLU. (#3635) (#3473) (#3666)
• [Common, PyTorch] Added fused_mla_kv_rope_mxfp8, which fuses multi-latent attention key/value RoPE with MXFP8 quantization for DeepSeek-V3 head dimensions on SM100 and later. (#3672)
• [Common, PyTorch] Added a fused RMSNorm forward kernel with MXFP8 output that writes compact or GEMM-swizzled scaling factors directly. (#3657)
• [Common] Improved low-precision quantization performance with 128x128-tile grouped MXFP8 requantization, one-CTA-per-tile MXFP8 scale swizzling, and single-pass row and column amax for single and grouped row-scaled NVFP4. (#3656) (#3658) (#3454) (#3467)
• [Common, PyTorch] Added NVFP4 Random Hadamard Transform support on SM120 and SM121 GPUs. (#3300)
• [Common, PyTorch] Added return_argmax=True to parallel_cross_entropy, returning the predicted token index without a second pass over the logits. (#3671)
• [PyTorch] Added zero_tensor_dummy_wgrad, which returns a storage-free dummy weight gradient for parameters using fuse_wgrad_accumulation, avoiding redundant additions when a weight is used more than once. (#3670)
• [Common, PyTorch] Reduced CPU overhead in Linear, LayerNormLinear, LayerNormMLP, and GroupedLinear by removing unnecessary tensor views and a redundant device-to-host split-size synchronization. (#3523) (#3554)
• [JAX] Optimized the MoE block with fused cuDNN Frontend MXFP8 grouped GEMM and activation kernels, quantization before FSDP all-gather, and user-controlled checkpoint names on block intermediates. (#3638) (#3589)
• [Build] Reduced source-build time by compiling NVFP4 quantization kernels with NVRTC. (#3301)
• [Build] Made aarch64 Transformer Engine libraries compatible with BOLT post-link optimization. (#3383) (#3627)
• [Docs] Added a datatype and hardware support matrix and a Mixture of Experts guide with PyTorch and JAX examples. (#3195) (#3494)
• [JAX] Added an opt-in fused Expert Parallelism prepare-and-dispatch path through NVTE_EP_FUSED_PREPARE_DISPATCH=1. (#3674)

Fixed Issues

• [Common] Fixed shared-memory reuse races and incomplete TMA store synchronization in FP8, MXFP8, NVFP4, and Hadamard-fusion quantization kernels. (#3417)
• [Common] Fixed incorrect column-wise scales in grouped MXFP8 quantization with varying last dimensions and GEMM-swizzled scales. (#3599)
• [Common] Fixed kernel launch failures on GPUs with smaller per-block shared-memory or register limits by falling back from gated-activation TMA kernels and bounding multi-tensor swizzle kernel registers. (#3600) (#3622)
• [Common] Fixed silent out-of-bounds bias reads in non-FP8 GEMMs by rejecting bias tensors whose dtype differs from the output dtype. (#3565)
• [Common] Fixed integer overflow in fused MoE auxiliary-loss kernels for inputs larger than INT_MAX elements. (#3569)
• [Common] Fixed loss of small sigmoid scores in biased fused top-k routing. (#3637)
• [Common, PyTorch] Fixed THD context-parallel kernel launch failures with large packed-document counts by keeping sequence-boundary tables within shared-memory limits. (#3698)
• [Common] Fixed segmentation faults at process exit after NVRTC kernels or CUDA graph-captured cuDNN fused attention were used. (#3733)
• [Common] Disabled FusedAttention for THD training with learnable softmax on cuDNN versions earlier than 9.26 to avoid incorrect dSink results and out-of-bounds reads. (#3470)
• [Common] Fixed Float8CurrentScaling and Float8BlockScaling ignoring the use_power_2_scales and use_f32_scales constructor arguments. (#3605)
• [Build, Common] Fixed source builds for Blackwell targets with CUDA 12.8. (#3602)
• [PyTorch] Fixed attention backend selection so a broken FlashAttention 2 or 3 installation no longer breaks import transformer_engine.pytorch, and FlashAttention 4 is skipped for deterministic backward on SM12x. (#3553) (#3595)
• [PyTorch] Made FlashAttention 4 backward on Hopper fail with an actionable error instead of hanging when quack-kernels 0.6.0 through 0.6.3 is installed with CUTLASS DSL 4.6.2 or later. (#3762)
• [PyTorch] Fixed NVFP4 torch.compile failures for non-matrix Linear inputs and incorrect gradients from compiled Linear operations with bias. (#3628)
• [PyTorch] Fixed ONNX export of MXFP8 block quantization for TensorRT 11.4, FP8 current-scaling quantization output metadata, and the registered layernorm operator. (#3592) (#3659)
• [PyTorch] Fixed deferred initialization of fusible operations by adding reset_parameters() to te.ops.Linear, te.ops.Bias, and te.ops.Sequential. (#3327)
• [PyTorch] Fixed GroupedLinear distributed weights on the grouped-tensor path, including backward main_grad handling and FP8 requantization of re-gathered weights. (#3517)
• [PyTorch] Fixed NVFP4 weight views in the single-group fused grouped MLP. (#3662)
• [PyTorch] Fixed GroupedTensor gradient propagation through autograd views and stream tracking of its backing tensors. (#3701)
• [PyTorch] Fixed LayerNormLinear silently disabling Userbuffers all-gather overlap when the forward pass runs under torch.no_grad(). (#3546)
• [PyTorch] Fixed FusedSGD.zero_grad(set_to_none=False) setting gradients to None instead of zeroing them. (#3549)
• [PyTorch] Fixed NCCL-EP bootstrap and Newton-Schulz with PyTorch's ProcessGroupNCCL2 backend by borrowing ...

Transformer Engine v2.20.2 Release Notes

Key Features and Enhancements

• [Common, JAX, PyTorch] Added an experimental CuTeDSL backend for MXFP8 quantization, enabled at build time with NVTE_WITH_CUTEDSL=1 and selected at runtime with NVTE_ENABLE_CUTEDSL_BACKEND=1, with automatic fallback to CUDA C++ kernels. (#3137)
• [Common, PyTorch] Added fused grouped MXFP8 paths for weighted SwiGLU and clamped-SwiGLU activation with columnwise quantization and for in-place requantization. (#3315) (#3359)
• [Common] Improved low-precision quantization performance with wider NVFP4 stochastic-rounding instructions, columnwise MXFP8 dBias reduction, register-resident MXFP8 cast kernels, and removal of an unused MXFP8 tile fetch. (#3357) (#3439) (#3459) (#3500)
• [Common, PyTorch] Added opt-in Quantile Balancing histogram accumulation to the fused sigmoid router, including two-kernel and fused-atomic modes. (#3395)
• [Common, PyTorch] Added native, row-scaled, and feature-gated fused grouped-MLP support for the SiTU-GLU activation. (#3402)
• [PyTorch] Enabled torch.compile(fullgraph=True) for Linear with supported FP8, MXFP8, and NVFP4 recipes and adopted current opaque custom-class registration where available. (#3053) (#3495)
• [PyTorch] Added per-sequence mask, window, and diagonal-alignment policies to packed THD attention, with automatic padded or grouped dispatch. (#3274)
• [Common, PyTorch] Added a CUDA graph-capturable fused Expert Parallelism prepare-and-dispatch path and reduced eager dispatch/combine CPU overhead. (#3282) (#3380)
• [PyTorch] Added NVTE_PLUGIN support for loading alternate PyTorch backend implementations and overriding FlashAttention and attention-backend selection. (#3401)
• [PyTorch] Reduced grouped-layer CPU overhead by avoiding redundant quantizer validation, reusing grouped-MLP fusion plans, and eliminating repeated device discovery. (#3405) (#3410)
• [PyTorch] Added paged-stashing metadata for saved GroupedLinear activation storage, including unquantized BF16 and FP16 rowwise tensors. (#3423)
• [PyTorch] Reduced CUDA graph memory retention by releasing warmup and capture outputs when consumed and clearing per-callable replay state on reset. (#3427)
• [PyTorch] Added ScaledTanhSReLU, a tanh-soft-clamped squared ReLU with per-row scaling and feature-gated cuDNN fused grouped-MLP support. (#3463)
• [PyTorch] Added experimental cuDNN-backed GatedDeltaProductAttention. (#3539)
• [JAX] Added Hopper BF16 grouped GEMM support. (#3083)
• [JAX] Added sqrtsoftplus scoring to fused MoE routing. (#3448)
• [JAX] Enabled Expert Parallelism to borrow XLA's NCCL communicator when supported, added single-process multi-device execution, and defaulted bootstrap to the AllGather-scan metadata path with fused scan-dispatch support. (#3452) (#3522) (#3561)
• [Build] Improved source-build dependency discovery for CUDA, NVCC, headers, and libraries installed through Python packages. (#3251)
• [Build] Reduced highly parallel source-build time by splitting grouped activation CUDA translation units. (#3430)

Fixed Issues

• [Common] Fixed illegal memory accesses and incorrect results in TMA-based quantization by preserving shared-memory pointer provenance and correcting Grouped MXFP8 work mapping and descriptor synchronization. (#3482) (#3483)
• [Common] Fixed BF16 GEMM precision loss on L40 by excluding cuBLASLt split-K algorithms that store partial reductions in BF16. (#3440)
• [Common, JAX] Fixed race conditions between zeroing and routed gradient stores in Triton permutation kernels. (#3581)
• [PyTorch] Fixed fused grouped-MLP correctness for E5M2 backward gradients, padded single-group SReLU inputs, and unsupported GEGLU or SiTU-GLU configurations on Rubin. (#3352) (#3485) (#3567)
• [PyTorch] Fixed cuDNN grouped-MLP weight-gradient compilation, CUDA graph capture, and FC1/FC2 corruption by retaining separate persistent descriptor workspaces. (#3579)
• [PyTorch] Fixed quantized parameter integration by supporting MXFP8 master-weight casts on ranks with empty shards and preserving quantized-tensor subclasses across detach() and C++ bindings. (#3348) (#3393)
• [PyTorch] Fixed delayed-scaling FP8 activation-recompute metadata pairing for evaluation mode, mode transitions, nested checkpoints, and checkpoint entry under disabled autograd. (#3394)
• [PyTorch] Fixed delayed weight-gradient computation when an external consumer clears an eagerly computed FP8 bias gradient before backward_dw(). (#3400)
• [PyTorch] Fixed FusedAdam with different parameter and gradient dtypes in capturable mode and prevented tensor-handle pool exhaustion for highly fragmented parameter groups. (#3414) (#3418)
• [PyTorch] Fixed Quantile Balancing router capture and replay with mutable bin bounds through an explicit trusted-producer validation handshake. (#3426)
• [PyTorch] Fixed non-fused and BF16 activation paths from failing when activation_recompute_in_mlp=True; unsupported paths now warn and proceed without recomputation. (#3436)
• [PyTorch] Fixed NCCL-EP zero-copy MXFP8 scale offsets for symmetric-memory windows and made offset resolution compatible with older and newer PyTorch APIs. (#3447) (#3466)
• [PyTorch] Made fused grouped-MLP dprob bit-exact when deterministic algorithms are requested and a supporting cuDNN Frontend is installed. (#3407)
• [PyTorch] Replaced deprecated TorchScript fallbacks with eager execution when torch.compile is disabled or unsupported. (#3502)
• [JAX] Fixed fused-attention backward output-gradient sharding, rotated THD metadata handling, and compound DP/FSDP test-data placement. (#3516)
• [PyTorch] Fixed NVFP4 stochastic-quantization corruption and invalid memory reads by retaining the owning RNG tensors through dispatch. (#3533)
• [Common, PyTorch] Fixed distributed Newton-Schulz numerical correctness, workspace sizing and lifetime, and validation of unsupported matrix dimensions. (#3541)
• [Build, PyTorch] Fixed source builds to gate NCCL-EP consistently on targeted GPU architectures and available PyTorch symmetric-memory support. (#3396)
• [Common, JAX, PyTorch] Fixed NCCL-EP synchronization, timeout, and shared-workspace failures and added validated compatibility for builds using NCCL 2.30.5–2.30.x with newer NCCL runtimes through NCCL_DEV_API_JIT. (#3619)

Breaking Changes in This Release

• [Common, JAX, PyTorch] Raised the minimum supported cuDNN version from 9.3 to 9.12.0, added configure-time, build-time, import-time, and runtime enforcement, and deprecated the now-obsolete experimental nvte_get_runtime_num_segments. (#3236)
• [Common, PyTorch] Added MXFP8 support to cuBLASMp communication-overlap GEMMs and raised the minimum cuBLASMp version for NVTE_WITH_CUBLASMP=1 from 0.8.0 to 0.8.1...

Transformer Engine v2.19 Release Notes

Key Features and Enhancements

• [Common] Added Rubin SM107a support when building with CUDA 13.4 or later. (#3275)
• [PyTorch] Added experimental per-direction hybrid quantization through CustomRecipe, allowing rowwise and columnwise paths to use different quantizers. (#2817)
• [Common, PyTorch] Added opt-in 2D block scaling for MXFP8 weights through MXFP8BlockScaling(enable_2d_quantization=True). (#2634)
• [Common, PyTorch] Added MXFP8 quantization support to NCCL-EP dispatch forward and combine backward, with GroupedLinear and fused grouped MLP support for prequantized grouped inputs. (#3270) (#3244)
• [Common, PyTorch] Added cuDNN FP8 fused attention support for packed THD inputs and context-parallel execution; FP8 THD sink-softmax backward requires cuDNN 9.26 or later. (#2994) (#3441)
• [PyTorch] Added experimental cuDNN-backed GatedDeltaNetAttention for dense and THD inputs. (#3351)
• [PyTorch] Enabled torch.compile(fullgraph=True) for supported FlashAttention and the non-FP8 unfused attention backend. (#3153) (#3189) (#3201) (#3286)
• [PyTorch] Added per-callable capture_time_hooks to make_graphed_callables for non-capturable warmup and graph-construction work. (#2831)
• [PyTorch] Added forward and backward support to FusedMLAQUpProjRopeQuant for fused MLA Q up-projection, RoPE, and MXFP8 quantization. (#3303) (#3330)
• [PyTorch] Expanded THD context parallelism with hierarchical a2a+p2p, experimental no-load-balance AllGather, and FlashAttention 4 on p2p, all_gather, and a2a paths. (#3290) (#3221) (#3438) (#3149)
• [PyTorch] Enabled row-scaled NVFP4 weight-gradient computation in dense Linear and GroupedLinear layers. (#3206) (#3324)
• [Common, PyTorch] Reduced MXFP8 grouped MLP CPU overhead. (#3338) (#3267) (#3238) (#3293)
• [PyTorch] Added named extra-output channels to te.ops.Sequential, allowing later fusible operations to consume outputs from earlier operations. (#3320)
• [PyTorch] Added newton_schulz_tp for tensor-parallel Newton-Schulz orthogonalization of replicated or sharded matrices. (#2920)
• [PyTorch] Optimized the performance and reduced parallel_cross_entropy memory use. Added overwrite_input option to reduce the memory usage even further by reusing the input storage. (#3273)
• [Common, PyTorch] Added NVFP4 stochastic-rounding and RHT split-quantization support on SM120 and SM121 GPUs. (#3281) (#3265)
• [JAX] Added per-head maximum-logit outputs to fused_attn and Flax DotProductAttention. (#3112)
• [JAX] Optimized experimental MoE Block to improve E2E perf, MXFP8 quantization, overflow detection, reduced other overheads (#3354)
• [JAX] Expert Parallelism APIs with reduced receive capacity and overflow detection (#3277)
• [JAX] Added ep_bootstrap support for meshes with axes orthogonal to expert parallelism by deriving communicator domains from the active mesh. (#3226)
• [Build] Added CMake discovery of NCCL headers installed by nvidia-nccl-cu12 and nvidia-nccl-cu13. (#3227)
• [Build] Enabled NVRTC to locate CUDA headers in Python site-packages at runtime, improving JIT compilation in relocated and build-isolated installations. (#3252)
• [Build] Reduced Transformer Engine library, wheel, and container sizes by compressing embedded CUDA binaries. (#3465)

Fixed Issues

• [Common] Disabled an affected cuBLAS 13.7 grouped GEMM algorithm on B300 and Rubin GPUs to prevent silent data corruption. (#3475)
• [PyTorch] Fixed FlashAttention 2 local and post-release validation, FlashAttention 3 sliding-window context-parallel checks, and FlashAttention 4 fallback for unsupported GPUs, head dimensions, or CUTLASS DSL configurations. (#3356) (#3341)
• [PyTorch] Fixed FusedAdam checkpoint step counters for parameter groups that are empty on some ranks. (#3318)
• [PyTorch] Fixed delayed-scaling FP8 activation recomputation when checkpointed callables enable FP8 in an inner autocast context. (#3284)
• [PyTorch] Fixed CUDA graph replay parameter-gradient lifetime by cloning returned parameter gradients by default, preventing later replays from overwriting retained gradients. (#2937)
• [PyTorch] Fixed a possible first-call segfault during lazy extension initialization by acquiring the GIL before Python imports. (#3255)
• [PyTorch] Fixed THD fused RoPE launch scaling for high sequence counts in token-linear forward and backward kernels. (#3057)
• [JAX] Fixed GEMM partition inference for tuple-valued mesh-axis specifications, including nested contracting axes and DP+FSDP layouts. (#3429) (#3272)
• [Build] Fixed NCCL-EP JIT in installed packages by shipping the runtime library and required JIT headers and discovering matching NCCL headers at runtime. (#3434)
• [Build] Fixed NCCL-EP source builds to respect MAX_JOBS and use a valid unlimited-parallel make invocation when no limit is set. (#3138) (#3243)

Breaking Changes in This Release

• [PyTorch] Changed mhc_fused_projection so its H output is always FP32 instead of matching the input dtype; cast it explicitly where BF16 is required. (#2978)
• [PyTorch] Changed the NCCL-EP Python API to support eager dispatch and drop-on-overflow: pass num_topk to ep_bootstrap, update positional ep_bootstrap and EpBuffer calls, and replace EpBuffer.token_counts with EpBuffer.tokens_per_expert. (#3229)

Deprecated Features

• [PyTorch] Deprecated the unused is_cg_capturable argument to parallel_cross_entropy; the operation is always CUDA graph capturable. (#3455)
• [PyTorch] Deprecated selecting GroupedLinear's grouped-tensor path with NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM; pass use_grouped_tensor=True or False to GroupedLinear instead. (#3224)
• [PyTorch] Deprecated implicit pointer-based layout detection when DotProductAttention receives Q/K/V views of packed buffers; pass the packed buffer through qkv_layer or kv_layer and set qkv_interleave_dim instead. (#3200)

Known Issues in This Release

• [PyTorch] GDN attention may generate a NaN in the backward pass when used with head_dim==32. The fix has been merged into nvidia-cudnn-frontend and will be resolved by the next release.

Transformer Engine v2.18 Release Notes

Key Features and Enhancements

• [Common] Added fused BF16 forward/backward kernels for scaled variants of SwiGLU, ClampedSwiGLU, and SReLU. (#3132)
• [Common] Reduced binary size and build time by migrating fused softmax, LayerNorm, and RMSNorm kernels to NVRTC. (#3156)
• [Common] Improved the performance of variable-length cuDNN fused attention forward passes by passing sequence-length metadata directly to cuDNN. (#3186)
• [Common, PyTorch] Added dense topk_indices output support to the fused router for TE Expert Parallelism. (#3129)
• [Common, PyTorch] Added FP8 current-scaling grouped quantization with support for varying tensor dimensions. (#3114)
• [Common, PyTorch] Made FP8 current scaling in GroupedLinear compatible with CUDA graph capture and improved cuBLASLt grouped GEMM algorithm selection. (#3143)
• [Common, PyTorch] Added CUDA graph capture support for THD attention. (#2898)
• [Common, JAX, PyTorch] Enabled cuDNN fused attention backward for attention head dimension 256 on Blackwell GPUs with compute capability 10.x. (#3056)
• [PyTorch] Enabled FP8 block scaling in GroupedLinear and fusible ops. (#3171) (#3242) (#3135)
• [PyTorch] Reduced scale-factor swizzling overhead for cached MXFP8 and NVFP4 weights. (#3093)
• [PyTorch] Added THD format support for AllGather-based context parallelism. (#2829)
• [PyTorch] Added zero-copy support to TE Expert Parallelism. (#3187)
• [PyTorch] Enabled NVFP4 RHT amax fusion in the grouped MLP using SRelu activation. (#3133)
• [PyTorch] Added NVTE_FLASH_ATTN_V2, NVTE_FLASH_ATTN_V3, and NVTE_FLASH_ATTN_V4 environment variables to control each FlashAttention version independently. (#3204)
• [JAX] Integrated the TE EP backend into the experimental moe() VJP API and experimental Flax MoEBlock. (#3116)
• [Docs] Improved the documentation on attention (#3142) (#3162)
• [Build] Replaced the cuDNN Frontend submodule with the nvidia-cudnn-frontend python package for source builds. (#3169) (#3278)

Fixed Issues

• [Common, PyTorch] Fixed grouped GEMM failures with FP8 per-tensor scaling on Hopper when using cuBLAS versions earlier than 13.5 by falling back to the legacy path. (#3181)
• [Common, PyTorch] Fixed CUDA graph failures with cached NVFP4 weights by producing GEMM-swizzled scaling factors during quantization when supported and disabling pre-swizzling otherwise. (#3190) (#3232)
• [PyTorch] Fixed gradient propagation for quantized tensors across autograd boundaries and made FP8 block-scaling dequantization autograd-aware. (#3172)
• [PyTorch] Fixed NVTE_BACKWARD_OVERRIDE=dequantized to use dequantized forward operands in backward. (#3141)
• [PyTorch] Fixed precision loss in the fused cross-entropy backward pass. (#3193)
• [PyTorch] Fixed process aborts during communication+GEMM overlap and NCCL-EP initialization caused by incorrect GIL handling. (#3203)
• [PyTorch] Fixed skipped updates in FusedAdam and FusedSGD when parameter groups contain empty tensors. (#3212)
• [PyTorch] Fixed NCCL communicator initialization when creating a cuSOLVERMp context. (#3240)
• [PyTorch] Fixed quantized tensor __repr__ failures when data cannot be materialized, including under torch.compile. (#3146)
• [JAX] Fixed the distributed MoE auxiliary loss backward pass by preserving the scalar cotangent. (#3237)
• [Build] Fixed a CUDA 13.0 source build failure in grouped FP8 quantization code. (#3194)

Known Issues in This Release

• There is a known compatibity issue between the FlashAttention v4, CuTeDSL and CUDNN Frontend pip packages, which could produce runtime errors like ImportError: cannot import name 'block_copy' from 'cutlass.utils'. A supported stable combination:
  • flash-attn-4==4.0.0b11
  • nvidia-cutlass-dsl[cu13]==4.4.2
  • nvidia-cudnn-frontend==1.26.0

Breaking Changes in This Release

There are no breaking changes in this release.

Deprecated Features

There are no deprecated features in this release.

Transformer Engine v2.17.1 Patch Release Notes

Fixed Issues

• [PyTorch] Fixed an issue where the "extra state" of a checkpoint was unpickled even if the quantization recipe is stateless. To enhance security, add an explicit guard via environment variable in order to load the "extra state" portion of the checkpoints for stateful recipes. (#3123).

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.

Transformer Engine v2.16.1 Patch Release Notes

Fixed Issues

• [PyTorch] Fixed an issue where the "extra state" of a checkpoint was unpickled even if the quantization recipe is stateless. To enhance security, add an explicit guard via environment variable in order to load the "extra state" portion of the checkpoints for stateful recipes. (#3123).

Transformer Engine v2.16 Release Notes

Key Features and Enhancements

• [Common] Improved the performance of the split-overlap reduce-scatter GEMMs. (#2056)
• [Common] Improved the fused MoE auxiliary loss kernel performance for models with a large number of experts. (#2758)
• [Common] Optimized MXFP8 and NVFP4 dequantize kernels for improved performance. (#2865)
• [Common] Improved performance of the MXFP8 quantization kernels. (#2958)
• [PyTorch] Added pad_between_seqs support for non-CP and CP (A2A and P2P) with FA3 + THD (varlen) attention. (#2596)
• [PyTorch] Added role-based custom quantization control, enabling recipes to target specific modules and tensor types. (#2620)
• [PyTorch] Added end-to-end Mixtral MoE examples showing TE GroupedLinear integration with HuggingFace models for BF16 and FP8 training. (#2642)
• [PyTorch] Increased performance of the CPU activation offloading path in some cases (#2793)
• [PyTorch] Reduced the CPU overhead in the GroupedLinear module and operation (#2900) (#2957) (#2666)
• [PyTorch] Added CUDA Graph capture support for GroupedLinear and grouped MoE operations on supported configurations. (#2923)
• [PyTorch] Added FlashAttention 4 support for attention head dimension 256. (#2932)
• [JAX] Improved MoE permutation kernel performance. (#2975)
• [JAX] Improved JAX tutorial documentation with updated examples and guidance. (#2976)
• [Common, PyTorch] Added bias and dbias support for GroupedLinear layers. (#2885)
• [Common, PyTorch] Added variable grouped swizzle support for flexible grouped tensor memory layouts. (#2914)
• [Common, PyTorch] Implemented a row-scaled NVFP4 forward propagation recipe. (#2931)
• [Common, PyTorch] Expanded grouped GEMM support with NVFP4 on Blackwell and FP8 block scaling on Hopper. (#2971)
• [Common, JAX] Added a top-k operation for faster MoE routing. (#2890)
• [Common, JAX] Enabled the cuDNN fused attention backend for no-mask bidirectional sliding-window attention. (#2961)

Fixed Issues

• [PyTorch] Fixed variable-length attention cache reuse across devices and inference/training modes. (#2728)
• [PyTorch] Fixed FSDP2 memory leaks for FP8 weight workspaces and transpose caches. (#2805)
• [PyTorch] Fixed TE fuser behavior in torch.no_grad() paths by avoiding invalid gradient-flag updates on non-leaf tensors. (#2919)
• [PyTorch] Fixed distributed checkpoint loading for FSDP2 for models initialized with QuantizedModelInit. (#2974)
• [Common, PyTorch] Fixed cuBLAS grouped GEMM when weight dimensions are not divisible by 128. (#2954)
• [Common, PyTorch] Fixed int32 overflow and -1 sentinel value handling in moe_permute. (#2907)
• [Common, PyTorch] Fixed context-parallel FlashAttention output handling when FA3 is installed without FA2.(#2825)
• [Common, PyTorch] Disabled RHT quantization fusion on unsupported GPU architectures to avoid launch failures. (#2968)
• [PyTorch] Fixed a crash coming from GroupedLinear weight-gradient allocation. (#3049)

Breaking Changes in This Release

• [Common, PyTorch] The original FP8 delayed-scaling fused attention path has been removed. FP8 attention now uses the current cuDNN-backed implementation. (#2959)
• [Common, PyTorch, JAX] Removed the legacy f16_max512 fused-attention backend. BF16/FP16 attention is routed through the maintained arbitrary-sequence backend, but explicit selections of the old backend must be updated. (#2949)

Deprecated Features

There are no deprecated features in this release.

Transformer Engine v2.15 Release Notes

Key Features and Enhancements

• [PyTorch] Added support for Flash Attention 4. (#2432)
• [PyTorch] Added support for MXFP8 attention. (#2719)
• [PyTorch] Added support for QGeGLU activation both in te.ops and the fused grouped MLP path using GEMM + activation fusion. (#2855)
• [PyTorch] Added support for per-token bias probability scaling both in te.ops and the fused grouped MLP path using GEMM + activation fusion. (#2864)
• [PyTorch] Added support for NVFP4 weight quantization in the fused Adam optimizer. (#2797)
• [PyTorch, Common] Added triton kernels to support mHC (Manifold-Constrained Hyper-Connections). (#2790)
• [PyTorch, Common] Added support for dequantizing MXFP8 grouped tensors. (#2722)
• [Common] Added support for unswizzling scaling factors. (#2837,#2732)
• [PyTorch] Added Newton–Schulz orthogonalization via cuSOLVERMp for distributed orthogonalization workloads. (#2706)
• [PyTorch] Added an NVTE_BACKWARD_OVERRIDE=high_precision|dequantized environment variable to control backward precision behavior. (#2644)
• [PyTorch] Added a feature to debug tools to allow tensor dumps before and after quantization for numerical debugging. (#2645)
• [PyTorch] Optimized FP8 block-scaling AllGather for FSDP2 to reduce communication overhead. (#2789)
• [PyTorch] Added an example demonstrating high-precision weight initialization with fully_shard. (#2785)
• [PyTorch] Expanded fused grouped MLP support via te.ops by lowering the weight dimension requirements to being divisible by 64 (previously 256). (#2856)
• [PyTorch] Added torch.compile support for the MoE permute utility functions. (#2686)
• [Common, PyTorch] Improved the performance of NVFP4 quantization by refactoring the amax compute kernel. (#2820)
• [JAX] Reduced THD seqlen and offset computation from O(T·T) memory down to O(T) for long sequences. (#2522)
• [JAX] Added MXFP8 grouped quantize + GEMM support. (#2763)

Fixed Issues

• [PyTorch] Fixed a numerical bug where stale columnwise weight data would be used for post-validation training steps. (#2929)
• [PyTorch] Fixed redundant memory usage when using NVFP4 parameters. (#2834)
• [JAX] Fixed the JAX extension build with NVTE_UB_WITH_MPI=1. (#2835)
• [Common] Fixed a numerical bug for the MoE fused router for large top-K and expert counts. (#2821)
• [Common] Fixed an illegal memory access in register_user_buffer_collective on Ampere (and older) GPUs when using user buffers for COMM-GEMM overlap. (#2859)
• [Build] Fixed a build crash when compiling from source with NVTE_CUDA_ARCHS=120. (#2832)

Known issues

• [PyTorch] When building a grouped MLP module via te.ops.Sequential in order to use the GEMM + activation fusion, the kernel may produce non-deterministic results in the single grouped-weight case (i.e., when the environment variable NVTE_GROUPED_LINEAR_SINGLE_PARAM and the corresponding module argument single_grouped_weight is set).
• [PyTorch] Enabling fused grouped MLP via te.ops requires cudnn-frontend library version 1.23.0. In case of issues please ensure that the right version of CuTeDSL is correctly installed:

python -m pip uninstall -y \
  cutlass \
  nvidia-cutlass \
  nvidia-cutlass-dsl \
  nvidia-cutlass-dsl-libs-base \
  nvidia-cutlass-dsl-libs-cu13 \
  nvidia-cudnn-frontend
python -m pip install -U pip setuptools wheel
python -m pip install --no-cache-dir "nvidia-cutlass-dsl[cu13]==4.4.1"
python -m pip install --no-cache-dir "nvidia-cudnn-frontend[cutedsl]==1.23.0"

Breaking Changes in This Release

There are no breaking changes in this release.

Deprecated Features

There are no deprecated features in this release.

Transformer Engine v2.14.1 Patch Release Notes

Fixed Issues

• [All] Fixed the issue where the MXFP8 quantization + dbias fusion could nondeterministically produce the wrong results (#2921).