[CI] Add workflow verifying build against minimum supported PyTorch and JAX

.github/workflows/minimum_versions.yml

@@ -0,0 +1,83 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+# Verify that TE builds and imports against the minimum supported framework versions.
+name: 'Minimum supported versions'
+on:
+  pull_request:
+  workflow_dispatch:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+env:
+  MIN_TORCH_VERSION: "2.8.0"
+  MIN_JAX_VERSION: "0.5.3"
+jobs:
+  min-pytorch:
+    name: 'Minimum supported PyTorch'
+    runs-on: ubuntu-latest
+    container:
+      image: nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04
+      options: --user root
+    steps:
+      - name: 'Dependencies'
+        run: |
+          apt-get update
+          apt-get install -y --allow-change-held-packages git python3.9 pip cudnn9-cuda-12 libnccl-dev libnccl2
+          pip install torch==${MIN_TORCH_VERSION}
+          pip install cmake pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging einops onnxscript "nvidia-cudnn-frontend>=1.25.0"
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: ccache
+        uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad
+      - name: 'Build'
+        run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation --no-deps . -v
+        env:
+          NVTE_FRAMEWORK: pytorch
+          # Single arch to keep memory/build size down; sm90 so the NCCL EP
+          # path (and its torch-version gate) is exercised
+          NVTE_CUDA_ARCHS: "90"
+          # Prefer apt cudnn9 headers over the pip nvidia-cudnn-cu12 copy
+          CUDNN_INCLUDE_PATH: /usr/include/x86_64-linux-gnu
+          CUDNN_LIBRARY_PATH: /usr/lib/x86_64-linux-gnu
+          # Full parallelism OOMs the 7GB runner (exit 137)
+          MAX_JOBS: 2
+          SCCACHE_GHA_ENABLED: "true"
+      - name: 'Sanity check'
+        # No GPU driver on the runner; the sm90 build links libcuda via NCCL EP
+        run: |
+          ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/lib/x86_64-linux-gnu/libcuda.so.1
+          python3 tests/pytorch/test_sanity_import.py
+  min-jax:
+    name: 'Minimum supported JAX'
+    runs-on: ubuntu-latest
+    container:
+      image: nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04
+      options: --user root
+    steps:
+      - name: 'Dependencies'
+        run: |
+          apt-get update
+          apt-get install -y git python3.9 pip cudnn9-cuda-12
+          pip install jax==${MIN_JAX_VERSION} "flax>=0.7.1"
+          pip install cmake pybind11[global] ninja packaging pydantic "nvidia-cudnn-frontend>=1.25.0"
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: ccache
+        uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad
+      - name: 'Build'
+        run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation --no-deps . -v
+        env:
+          NVTE_FRAMEWORK: jax
+          NVTE_CUDA_ARCHS: "70"
+          CUDNN_INCLUDE_PATH: /usr/include/x86_64-linux-gnu
+          CUDNN_LIBRARY_PATH: /usr/lib/x86_64-linux-gnu
+          MAX_JOBS: 2
+          SCCACHE_GHA_ENABLED: "true"
+      - name: 'Sanity check'
+        run: python3 tests/jax/test_sanity_import.py

build_tools/pytorch.py

@@ -15,6 +15,7 @@
     cuda_version,
     get_cuda_include_dirs,
     debug_build_enabled,
+    nccl_ep_enabled,
     setup_mpi_flags,
 )
 from typing import List
@@ -89,12 +90,27 @@ def setup_pytorch_extension(
     # Mirror the NCCL EP gate from setup.py / common CMake. When disabled, the
     # ep.cpp source no-ops at the #ifdef boundary; without the define it would
     # produce undefined references to nvte_ep_*.
-    if bool(int(os.getenv("NVTE_WITH_NCCL_EP", "1"))):
-        cxx_flags.append("-DNVTE_WITH_NCCL_EP")
-        # PyTorch's symm-mem headers gate the NCCL_HAS_SYMMEM_* feature macros on
-        # USE_NCCL. The EP extension shares the symm-mem NCCL comm with torch, so
-        # it needs those macros visible.
-        cxx_flags.append("-DUSE_NCCL")
+    if nccl_ep_enabled():
+        # ep.cpp additionally needs torch's NCCL symm-mem headers (torch >= 2.11).
+        from torch.utils.cpp_extension import include_paths
+
+        symm_mem_header = "torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp"
+        if any(os.path.exists(os.path.join(p, symm_mem_header)) for p in include_paths()):
+            cxx_flags.append("-DNVTE_WITH_NCCL_EP")
+            # PyTorch's symm-mem headers gate the NCCL_HAS_SYMMEM_* feature macros on
+            # USE_NCCL. The EP extension shares the symm-mem NCCL comm with torch, so
+            # it needs those macros visible.
+            cxx_flags.append("-DUSE_NCCL")
+        elif os.getenv("NVTE_WITH_NCCL_EP") == "1":
+            raise RuntimeError(
+                "NVTE_WITH_NCCL_EP=1 was set but the installed torch does not provide "
+                f"{symm_mem_header}. NCCL EP requires torch >= 2.11."
+            )
+        else:
+            print(
+                f"[NCCL EP] Installed torch does not provide {symm_mem_header} "
+                "(torch >= 2.11 required); skipping NCCL EP in the torch extension."
+            )
 
     library_dirs = []
     libraries = []