v2.18

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.

v2.17.1

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).

v2.17

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.

v2.16.1

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).

v2.16

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.

v2.15

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.

v2.14.1

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).

v2.14

Transformer Engine v2.14 Release Notes

Key Features and Enhancements

• [PyTorch] Added multiple CPU overhead optimizations across the framework integration to reduce per-step Python/host overhead. (#2559) (#2724)
• [C, PyTorch] Added BF16 and MXFP8 grouped GEMM support with on-device group sizes. (#2748) (#2669)
• [PyTorch] Added a fused GEMM + SwiGLU grouped MLP for MXFP8 to accelerate MoE forward/backward. (#2769)
• [PyTorch] Added support for a single-parameter GroupedLinear configuration, where the weights of all experts are stored in a single parameter, which reduces CPU overheads. (#2731)
• [PyTorch] Added backwards-compatible checkpoint support for the new single-parameter GroupedLinear. (#2761)
• [PyTorch] Extended the fused attention API to optionally return softmax Stats always and Max when return_max_logit=True, exposing more cuDNN intermediates to users. (#2677)
• [PyTorch] Enabled SM120 support for the fused attention path when cuDNN >= 9.18.1 is available. (#2693)
• [PyTorch] Added support for MXFP8BlockScaling and Float8BlockScaling quantized weight in FusedAdam. (#2753)
• [PyTorch] Added CUDA graph-compatible multi_tensor_scale_tensor API in the optimizer. (#2594)
• [PyTorch] Enabled CUDA Graph capture of modules with CPU offloading. (#2435)
• [PyTorch] Added support for non-FP32 params_dtype when using QK-normalization. (#2718)
• [PyTorch] Added precision debug-tools support for quantized model parameters. (#2141)
• [JAX] Added a JAX-side API to invoke the fused MoE router kernels. (#2711)
• [JAX] Integrated BF16 grouped GEMM with on-device group sizes. (#2680)
• [JAX] Added a Collective GEMM (CGEMM) implementation with FP8 and MXFP8 support. (#2740)
• [JAX] Added Shardy support to the Collective GEMM (CGEMM) path. (#2714)
• [JAX] Improved the performance of the permutation kernels for the JAX 0.8.0 and newer. (#2741)
• [C] Enabled the fused RMSNorm dLN + add backward path through cuDNN for faster fused-residual normalization. (#2778)
• [C] Added a grouped MXFP8 quantization kernel, including grouped dbias support. (#2738) (#2674)
• [C] Enabled dequantization from an MXFP8 tensor that only carries column-wise data. (#2712)
• [C/PyTorch] Improved the performance of the NVFP4 recipe by fusing row-cast / RHT / transpose / column-cast. (#2555)
• [C] Made the number of Philox rounds for stochastic rounding configurable. (#2751)
• [Documentation] Added a documentation page describing CPU offloading in Transformer Engine. (#2520)
• [Documentation] Updated the documentation to describe the current cuDNN sliding-window attention support. (#2624)
• [Documentation] Improved error messages across the C, PyTorch, and JAX layers. (#2705)
• [Documentation] Added a custom-feature tutorial for the precision debug tools. (#2216)
• [Documentation] Added documentation for the operator fuser API. (#2447)
• [PyTorch, Documentation] Added end-to-end examples for fused_adam, quantized_model_init, and FSDP2 usage. (#2698) (#2662)

Fixed Issues

• [PyTorch] FSDP2 / Megatron-FSDP / DCP (distributed checkpointing): when model parameters are DTensors, ensure optimizer states are also DTensors for correct sharded checkpoints. (#2795)
• [PyTorch] Fixed async DCP checkpointing for Float8Tensor parameters. (#2721)
• [PyTorch] Fixed the issue with cross_entropy_forward producing wrong answers for non-contiguous logits. (#2746)
• [PyTorch] Fixed the excessive memory usage issue when using operator fuser. (#2750)
• [PyTorch] Fixed a precision-debug-tools crash when tp_group=None. (#2733)
• [PyTorch] Fixed Flash Attention 3 API compatibility for the window-size parameters. (#2704)
• [PyTorch] Fixed the initialization of the learnable softmax_offset parameter in DotProductAttention to zero-initialization. (#2694)
• [PyTorch] Fixed the error with FP8 block scaling when sequence parallelism is enabled and local tensor dimensions are not divisible by 128. (#2637)
• [PyTorch] Added a clear error when constructing LayerNormLinear with row-wise tensor parallelism (an unsupported configuration). Previously this configuration would fail with the CUDA error (#2688)
• [JAX] Fixed the performance issue with THD/BSHD segment-position generation. (#2823)
• [JAX] Fixed the assertion error when using from_segment_ids_and_pos() with vmap. (#2692)
• [JAX] Fixed the performance issue for models using both FSDP and EP. (#2649)
• [JAX] Changed the dtype of the intermediate-result aval in fused_topk_and_score_function_fwd to fp32 to avoid precision loss. (#2752)
• [C] Fixed an incorrect MNNVL fabric-availability check that misreported support on some systems. (#2626)
• [C/PyTorch] Fixed score normalization in fused_score_for_moe_aux_loss when topk == 1. (#2720)
• [PyTorch] Fixed the possible precision loss when copying from the quantized tensor to the high precision tensor. (#2120, #2673)

Breaking Changes in This Release

• [JAX] GSPMD partitioning rules are no longer tested and will now warn on use; users on JAX with GSPMD should migrate to Shardy. (#2702)

Deprecated Features

There are no deprecated features in this release.

v2.13

Transformer Engine v2.13 Release Notes

Key Features and Enhancements

• Added detailed documentation for low precision training with Transformer Engine, covering FP8, MXFP8, NVFP4, and other quantization recipes with examples for both PyTorch and JAX. (#2343).
• [Build] Added NVTE_BUILD_USE_NVIDIA_WHEELS environment variable to allow building TE using CUDA headers from PyPI NVIDIA wheels instead of a system CUDA installation. (#2623)
• [C] Enabled deterministic FP8 fused attention on Blackwell (SM100) GPUs. (#2621)
• [C] Updated cuBLASMp integration to version 0.8.0, replacing the nvshmem dependency with NCCL-based symmetric memory. (#2661)
• [C] Added MXFP8 quantization kernels for grouped tensors used in MoE, with fused scale-factor swizzling for improved performance. (#2586, #2630)
• [C] Added NVFP4 quantization kernels for grouped tensors used in MoE models. (#2655)
• [C] Reduced cuDNN graph recompilations in THD fused attention by rounding large batch sizes to 512-element increments. (#2653)
• [C] Added sqrtsoftplus scoring function to the fused MoE router and improved router kernel performance on Blackwell GPUs. (#2633, #2683)
• [PyTorch] Introduced GroupedTensor, enabling MoE expert weights to be stored as a single contiguous allocation while remaining individually addressable. (#2654)
• [PyTorch] Added fusible GroupedLinear and ScaledSwiGLU ops for building fully fused MoE grouped MLP pipelines. (#2664)
• [PyTorch] Added register_forward_fusion and register_backward_fusion APIs, allowing users to define and register custom operator fusion patterns. (#2597)
• [PyTorch] Added get_backward_dw_params API to TE modules, fixing weight gradient hook management when using wgrad CUDA Graphs with Megatron-LM. (#2614)
• [PyTorch] Fixed fused attention bias dimension handling and extended dbias support to additional bias shapes (b1ss, bhss, 11ss, 111s). (#2537)
• [PyTorch] Reduced peak memory usage in fused Adam optimizer by fusing BF16 momentum scaling directly into CUDA kernels, also enabling CUDA Graph capture for this path. (#2632)
• [PyTorch] Added the sigmoid-gated GLU activation (activation="glu") to LayerNormMLP and TransformerLayer. (#2656)
• [PyTorch] Extended debug statistics tracking to NVFP4 quantization (underflow and MSE metrics), and gracefully skipped stat logging for layers not using quantization. (#2296, #2652)
• [PyTorch] Fixed CUDA Graph capture for Megatron-Core vision encoder models. (#2657)
• [JAX] Added experimental inspect_array debugging utility for dumping tensor snapshots during multi-GPU execution. (#2651)
• [JAX] Fixed MoE permutation to correctly mask padding tokens and handle tensor sizes under expert parallelism. (#2672)
• [JAX] MoE permutation now always returns tokens_per_expert, required for ragged all-to-all communication in expert parallelism. (#2613)

Fixed Issues

• [C] Fixed incorrect results from the exp2f_rcp fast-math helper when inputs are NaN or have biased exponent 254. (#2647)
• [C] Fixed a race condition in Randomized Hadamard Transform amax kernels where a missing memory fence could cause incorrect amax values. (#2695)
• [PyTorch] Fixed the TE Llama example to work with HuggingFace Transformers 4.57+, which changed decoder layer output conventions. (#2572)
• [Build] Fixed TypeError during build when NCCL is installed from PyPI as a namespace package without a __file__ attribute. (#2580)
• [Build] Fixed ModuleNotFoundError when installing from cached source distributions (e.g., via uv) by including build_tools in MANIFEST.in. (#2684)

Breaking Changes in This Release

• [C] Removed the deprecated packed fused attention C APIs (nvte_fused_attn_{fwd,bwd}_{qkvpacked,kvpacked}); users must migrate to the non-packed API variants. (#2696)
• Versions of cuBLASMp prior to 0.8.0 are no longer supported.

Deprecated Features

No features deprecated in this release.

v2.12

Transformer Engine v2.12 Release Notes

Key Features and Enhancements

• Made miscellaneous improvements and fixes to the documentation.
• [C] Improved performance of NVFP4 quantization kernels. (#2412)
• [C] Documented environment variables. (#2552)
• [PyTorch] Added fused permute+pad and unpermute+unpad operations for FP8 optimization. (#1921)
• [PyTorch] Improved the performance in CPU-limited scenarios.
• [PyTorch] Added support for Sliding Window Attention (left, right) with fused attention. (#2477)
• [PyTorch] Improved the performance of MXFP8 and NVFP4 by fusing the swizzling into the quantization (#2486)
• [PyTorch] Added cudagraph support for activation recomputation. (#2518)
• [JAX] Added a tutorial for integrating TE/JAX quantization into existing frameworks. (#2423)
• [JAX] Added custom partitioning for permutation primitives. (#2591)

Fixed Issues

• [C] Fixed SM120 compilation with CUDA 12. (#2482)
• [C] Fixed overflow in padding and unpadding kernels. (#2548)
• [C] Fixed a numerical issue in sort_chunks_by_index. (#2566)
• [C] Fixed a numerical issue in swizzling blockwise E8 scales. (#2589)
• [PyTorch] Fixed an AttributeError issue when checkpointing the model with MXFP8 parameters. (#2427)
• [PyTorch] Fixed cross-entropy loss calculation when some tokens are ignored. (#2476)
• [PyTorch] Fixed Float8Tensor.contiguous autograd support. (#2533)
• [PyTorch] Fixed multiple CPU offloading issues. (#2535)
• [PyTorch] Fixed uninitialized permuted_scale values. (#2547)
• [PyTorch] Fixed FP8 quantization for the second MLP in LayerNormMLP. (#2577)
• [PyTorch] Fixed ONNX tests and added FP8 attention export support. (#2598)
• [JAX] Removed unused TE DPA dtype handling to improve cuDNN backend dtype detection. (#2485)
• [JAX] Fixed segment-position calculation from segment IDs in SequenceDescriptor class. (#2523)
• [JAX] Fixed bugs in permutation custom partitioning. (#2617)
• [JAX] Fixed issue in encoder and MNIST examples due to dataset path moving. (#2625)

Breaking Changes in This Release

No breaking changes in this release.

Deprecated Features

No features deprecated in this release.