IFU v2.14.dev0

.github/workflows/build.yml

@@ -7,6 +7,10 @@ name: 'Build'
 on:
   pull_request:
   workflow_dispatch:
+concurrency:
+  # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes)
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
 jobs:
   core:
     name: 'Core'

.github/workflows/deploy_nightly_docs.yml

@@ -7,6 +7,7 @@ name: Deploy nightly docs
 on:
   push:
     branches: [ "main" ]
+  workflow_dispatch:
 jobs:
   build:
     uses: ./.github/workflows/docs.yml
@@ -21,9 +22,8 @@ jobs:
             name: "te_docs"
             path: "html"
       - name: Prepare for pages
-        uses: actions/upload-pages-artifact@v1.0.7
+        uses: actions/upload-pages-artifact@v3
         with:
-          name: github-pages
           path: "html"
   deploy:
     needs: prepare
@@ -36,4 +36,5 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Deploy
-      uses: actions/deploy-pages@v2.0.0
+      id: deployment
+      uses: actions/deploy-pages@v4

.github/workflows/docs.yml

@@ -8,6 +8,10 @@ on:
   pull_request:
   workflow_dispatch:
   workflow_call:
+concurrency:
+  # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes)
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
 jobs:
   build_docs:
     name: 'Build'

.github/workflows/lint.yml

@@ -7,6 +7,10 @@ name: 'Lint'
 on:
   pull_request:
   workflow_dispatch:
+concurrency:
+  # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes)
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
 jobs:
   pytorch_cpplint:
     name: 'PyTorch C++'

.github/workflows/trigger-ci.yml

@@ -58,6 +58,8 @@ jobs:
            || github.actor == 'vthumbe1503'
            || github.actor == 'shengfangd'
            || github.actor == 'kainzhong'
+           || github.actor == 'cspades'
+           || github.actor == 'jomitchellnv'
          )
     steps:
       - name: Check if comment is issued by authorized person

.gitignore

@@ -58,4 +58,8 @@ artifacts/
 **/times.csv
 transformer_engine/build_info.txt
 transformer_engine/common/util/hip_nvml.*
-*.DS_Store
+.DS_Store
+.rsync-filter
+.codex/
+.cline_storage/
+.claude/

.pre-commit-config.yaml

@@ -43,4 +43,4 @@ repos:
     rev: c75aca72f4e85c6e47252139e8695f1c8b5f9ae3
     hooks:
       - id: vermin
-        args: ['-t=3.10', '--violations']
+        args: ['-t=3.10-', '--violations']

MANIFEST.in

@@ -1 +1,2 @@
 recursive-include transformer_engine/common/include *.*
+recursive-include build_tools *.py *.txt

README.rst

@@ -458,7 +458,7 @@ Flax
       for _ in range(10):
         loss, (param_grads, other_grads) = fwd_bwd_fn(params, other_variables, inp)
 
-For a more comprehensive tutorial, check out our `Quickstart Notebook <https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/quickstart.ipynb>`_.
+For a more comprehensive tutorial, check out our `Getting Started Guide <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/getting_started.html>`_.
 
 .. overview-end-marker-do-not-remove
 
@@ -496,15 +496,22 @@ For example to use the NGC PyTorch container interactively,
 
 .. code-block:: bash
 
-    docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.08-py3
+    docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.01-py3
 
 For example to use the NGC JAX container interactively,
 
 .. code-block:: bash
 
-    docker run --gpus all -it --rm nvcr.io/nvidia/jax:25.08-py3
+    docker run --gpus all -it --rm nvcr.io/nvidia/jax:26.01-py3
 
-Where 25.08 (corresponding to August 2025 release) is the container version.
+Where 26.01 (corresponding to January 2026 release) is the container version.
+
+We recommend updating to the latest NGC container available here:
+
+* https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
+* https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax
+
+If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. If you would like to use examples from TE main branch and are running into import errors, please try the latest pip package or building from source, although NGC containers are recommended for ease-of-use for most users.
 
 **Benefits of using NGC containers:**
 
@@ -628,6 +635,37 @@ Troubleshooting
          cd transformer_engine
          pip install -v -v -v --no-build-isolation .
 
+**Problems using UV or Virtual Environments:**
+
+1. **Import Error:**
+
+   * **Symptoms:** Cannot import ``transformer_engine``
+   * **Solution:** Ensure your UV environment is active and that you have used ``uv pip install --no-build-isolation <te_pypi_package_or_wheel_or_source_dir>`` instead of a regular pip install to your system environment.
+
+2. **cuDNN Sublibrary Loading Failed:**
+
+   * **Symptoms:** Errors at runtime with ``CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED``
+   * **Solution:** This can occur when TE is built against the container's system installation of cuDNN, but pip packages inside the virtual environment pull in pip packages for ``nvidia-cudnn-cu12/cu13``. To resolve this, when building TE from source please specify the following environment variables to point to the cuDNN in your virtual environment.
+
+
+     .. code-block:: bash
+
+        export CUDNN_PATH=$(pwd)/.venv/lib/python3.12/site-packages/nvidia/cudnn
+        export CUDNN_HOME=$CUDNN_PATH
+        export LD_LIBRARY_PATH=$CUDNN_PATH/lib:$LD_LIBRARY_PATH
+
+3. **Building Wheels:**
+
+   * **Symptoms:** Regular TE installs work correctly but UV wheel builds fail at runtime.
+   * **Solution:** Ensure that ``uv build --wheel --no-build-isolation -v`` is used during the wheel build as well as the pip installation of the wheel. Use ``-v`` for verbose output to verify that TE is not pulling in a mismatching version of PyTorch or JAX that differs from the UV environment's version.
+
+**JAX-specific Common Issues and Solutions:**
+
+1. **FFI Issues:**
+
+   * **Symptoms:** ``No registered implementation for custom call to <some_te_ffi> for platform CUDA``
+   * **Solution:** Ensure ``--no-build-isolation`` is used during installation. If pre-building wheels, ensure that the wheel is both built and installed with ``--no-build-isolation``. See "Problems using UV or Virtual Environments" above if using UV.
+
 .. troubleshooting-end-marker-do-not-remove
 
 Breaking Changes

benchmarks/attention/benchmark_attention_jax.py

@@ -13,13 +13,11 @@
 from transformer_engine.jax.attention import (
     AttnBiasType,
     AttnMaskType,
+    AttnSoftmaxType,
     QKVLayout,
 )
 from transformer_engine.jax import fp8_autocast
 
-# Needed in order to dump timings properly
-os.environ["XLA_FLAGS"]="--xla_gpu_graph_level=0"
-
 # Add test_fused_attn to the sys path 
 tests_path = os.path.abspath(
     os.path.join(os.path.dirname(__file__), "../../tests/jax/")
@@ -108,12 +106,14 @@ def bench_forward(self, warmup, iters, timings_dir):
             jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
             jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding),
             jax.device_put(self.bias, self.bias_sharding),
+            jax.device_put(self.softmax_offset, self.softmax_offset_sharding),
             jax.device_put(self.sequence_desciptor, self.seq_desc_sharding),
             jax.device_put(self.dropout_rng, self.dropout_rng_sharding),
         ]
         kwargs = {
             "attn_bias_type": self.attn_bias_type,
             "attn_mask_type": self.attn_mask_type,
+            "softmax_type": self.softmax_type,
             "scaling_factor": self.scaling_factor,
             "dropout_probability": self.dropout_prob,
             "is_training": self.is_training,
@@ -122,6 +122,7 @@ def bench_forward(self, warmup, iters, timings_dir):
             "window_size": self.window_size,
             "context_parallel_strategy": self.cp_strategy,
             "context_parallel_causal_load_balanced": self.cp_load_balanced,
+            "stripe_size": self.stripe_size,
         }
 
         customcall_fused_dpa_jit = jax.jit(
@@ -132,6 +133,7 @@ def bench_forward(self, warmup, iters, timings_dir):
                 self.qkvo_sharding,
                 self.qkvo_sharding,
                 self.bias_sharding,
+                self.softmax_offset_sharding,
                 self.seq_desc_sharding,
                 self.dropout_rng_sharding,
             ],
@@ -180,12 +182,14 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs):
             jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
             jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding),
             jax.device_put(self.bias, self.bias_sharding),
+            jax.device_put(self.softmax_offset, self.softmax_offset_sharding),
             jax.device_put(self.sequence_desciptor, self.seq_desc_sharding),
             jax.device_put(self.dropout_rng, self.dropout_rng_sharding),
         ]
         kwargs = {
             "attn_bias_type": self.attn_bias_type,
             "attn_mask_type": self.attn_mask_type,
+            "softmax_type": self.softmax_type,
             "scaling_factor": self.scaling_factor,
             "dropout_probability": self.dropout_prob,
             "is_training": self.is_training,
@@ -194,6 +198,7 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs):
             "window_size": self.window_size,
             "context_parallel_strategy": self.cp_strategy,
             "context_parallel_causal_load_balanced": self.cp_load_balanced,
+            "stripe_size": self.stripe_size,
         }
 
         # We can compute dBias only for the [1, h, s, s] layout
@@ -212,8 +217,8 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs):
         # Use FP16/BF16 to sum the results may cause overflow, use FP32 for the summation
         jitted_primitive = jax.jit(
             jax.value_and_grad(
-                lambda q, k, v, bias, *args: grad_func(
-                    customcall_fused_dpa, q, k, v, bias, *args, cp_reverse_out=True, **kwargs
+                lambda q, k, v, bias, softmax_offset, *args: grad_func(
+                    customcall_fused_dpa, q, k, v, bias, softmax_offset, *args, cp_reverse_out=True, **kwargs
                 ),
                 arg_nums,
             ),
@@ -222,6 +227,7 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs):
                 self.qkvo_sharding,
                 self.qkvo_sharding,
                 self.bias_sharding,
+                self.softmax_offset_sharding,
                 self.seq_desc_sharding,
                 self.dropout_rng_sharding,
             ),
@@ -270,11 +276,13 @@ def _filter_configs(configs):
             continue
 
         backend = FusedAttnHelper(
+            is_training,
             dtype,
             dtype,
             qkv_layout,
             attn_bias_type,
             attn_mask_type,
+            AttnSoftmaxType.VANILLA_SOFTMAX,
             dropout_prob,
             h_q, h_kv,
             s_q, s_kv,
@@ -348,8 +356,8 @@ def benchmark_dot_product_attention_profiler(args):
             d_qk, d_v,
             attn_bias_type,
             attn_mask_type,
+            AttnSoftmaxType.VANILLA_SOFTMAX,
             dropout_prob,
-            True,
             dtype,
             is_training,
             qkv_layout,

build_tools/VERSION.txt

@@ -1 +1 @@
-2.12.0.dev0
+2.14.0.dev0

build_tools/hipify/custom_map.json

@@ -11,6 +11,7 @@
         "__nv_fp8_e5m2" : "te_hip_fp8_e5m2",
         "__nv_fp8_e4m3" : "te_hip_fp8_e4m3",
         "cuda::getCurrentCUDAStream" : "hip::getCurrentHIPStreamMasqueradingAsCUDA",
+        "c10::cuda::" : "c10::hip::",
         "__nv_fp4_e2m1" : "__hip_fp4_e2m1",
         "__nv_fp4x2_e2m1" : "__hip_fp4x2_e2m1",
         "__nv_fp4x4_e2m1" : "__hip_fp4x4_e2m1",

build_tools/hipify/hipify.py

@@ -112,8 +112,9 @@ def get_hipified_sources(hipify_result: Union[str, dict], sources: Union[list[Un
 
     # Because hipify output_directory == project_directory
     # Original sources list may contain previous hipifying results that ends up with duplicated entries
-    # Keep unique entries only
-    hipified_sources = set()
+    # Keep unique entries only but preserve the order by using list + set
+    hipified_sources = list()
+    hipified_sources_set = set()
     for fname in sources:
         if not os.path.isabs(fname):
             fname = os.path.join(src_base_path, fname)
@@ -122,15 +123,17 @@ def get_hipified_sources(hipify_result: Union[str, dict], sources: Union[list[Un
             file_result = hipify_result[fname]
             if file_result['hipified_path'] is not None:
                 fname = hipify_result[fname]['hipified_path']
-        hipified_sources.add(os.path.relpath(fname, str(src_base_path)))
+        if fname not in hipified_sources_set:
+            hipified_sources_set.add(fname)
+            hipified_sources.append(os.path.relpath(fname, str(src_base_path)))
 
     if sources_fname is not None:
         with open(sources_fname, "w") as f:
             for fname in hipified_sources:
                 f.write(fname + "\n")
         return sources_fname
 
-    return list(hipified_sources)
+    return hipified_sources
 
 
 def hipify_sources(te_root: Union[Path, str], src_dir: Union[Path, str],

build_tools/jax.py

@@ -59,12 +59,7 @@ def xla_path() -> str:
     Throws FileNotFoundError if XLA source is not found."""
 
     try:
-        import jax
-        from packaging import version
-        if version.parse(jax.__version__) >= version.parse("0.5.0"):
-            from jax import ffi
-        else:
-            from jax.extend import ffi
+        from jax import ffi
     except ImportError:
         if os.getenv("XLA_HOME"):
             xla_home = Path(os.getenv("XLA_HOME"))

build_tools/utils.py

@@ -306,9 +306,10 @@ def nvcc_path() -> Tuple[str, str]:
 def get_cuda_include_dirs() -> Tuple[str, str]:
     """Returns the CUDA header directory."""
 
+    force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0")))
     # If cuda is installed via toolkit, all necessary headers
     # are bundled inside the top level cuda directory.
-    if cuda_toolkit_include_path() is not None:
+    if not force_wheels and cuda_toolkit_include_path() is not None:
         return [cuda_toolkit_include_path()]
 
     # Use pip wheels to include all headers.        
@@ -317,7 +318,10 @@ def get_cuda_include_dirs() -> Tuple[str, str]:
     except ModuleNotFoundError as e:
         raise RuntimeError("CUDA not found.")
 
-    cuda_root = Path(nvidia.__file__).parent
+    if nvidia.__file__ is not None:
+        cuda_root = Path(nvidia.__file__).parent
+    else:
+        cuda_root = Path(nvidia.__path__[0])  # namespace
     return [
         subdir / "include"
         for subdir in cuda_root.iterdir()

ci/pytorch.sh

@@ -12,7 +12,7 @@ TEST_DIR=${TE_PATH}tests/pytorch
 #: ${TEST_WORKERS:=4}
 
 install_prerequisites() {
-    pip install 'numpy>=1.22.4' pandas
+    pip install 'numpy>=1.22.4' pandas safetensors
     rc=$?
     if [ $rc -ne 0 ]; then
         script_error "Failed to install test prerequisites"
@@ -104,8 +104,11 @@ run_test_config_mgpu(){
     run_default_fa 2 distributed/test_numerics.py
     run_default_fa 1 distributed/test_torch_fsdp2.py
     run_default_fa 2 distributed/test_torch_fsdp2_fp8.py
-    run_default_fa_lbl "flash" 3 attention/test_attention_with_cp.py -k "with_flash"
-    run_default_fa_lbl "fused" 2 attention/test_attention_with_cp.py -k "with_fused"
+    if [ $_fus_attn = ck ]; then
+        run 2 attention/test_attention_with_cp.py -k "with_fused"
+    elif [ $_fus_attn = flash ]; then
+        run 3 attention/test_attention_with_cp.py -k "with_flash"
+    fi
 }
 
 run_benchmark() {