From 02190d7fed3422d80e1db70eade26294c16baf1f Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 21 Oct 2026 00:44:02 +0000
Subject: [PATCH 1/6] Cache Triton kernels across CI runs

Every GPU and example container starts with an empty Triton cache, so each
job recompiles the same kernels (GDN, Mamba, FLA, the NVFP4 sweep, vLLM engine
start-up).

cache-extensions now also restores the newest Triton cache of the lane. Triton
entries are content-addressed and only ever added, so the restore uses
restore-keys on a run-unique key rather than an exact key. save-extensions saves
a new Triton cache only when a run compiled at least 20 entries more than it
restored, so a warm run saves nothing and a cache cannot be frozen incomplete by
a failed run. The extension cache keeps its exact key; its hit check moves from
the workflows' `if` into the action.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .github/actions/cache-extensions/action.yml | 23 ++++++++++++++--
 .github/actions/save-extensions/action.yml  | 29 ++++++++++++++++++---
 .github/workflows/_example_tests_runner.yml |  5 +++-
 .github/workflows/gpu_tests.yml             |  5 +++-
 4 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml
index a25dcb5fafc..e0eb96cf8a6 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/cache-extensions/action.yml
@@ -1,5 +1,5 @@
 name: Cache extensions
-description: Restore JIT-compiled extensions (e.g. CUDA extensions) cached by an earlier run. Pair with ./.github/actions/save-extensions.
+description: Restore JIT-compiled extensions (CUDA extensions, Triton kernels) cached by an earlier run. Pair with ./.github/actions/save-extensions.
 
 inputs:
   cache-key:
@@ -13,12 +13,20 @@ outputs:
   cache-key:
     description: The key to pass to ./.github/actions/save-extensions.
     value: ${{ steps.cache.outputs.cache-primary-key }}
+  triton-key:
+    description: The Triton key to pass to ./.github/actions/save-extensions.
+    value: ${{ steps.triton.outputs.cache-primary-key }}
+  triton-count:
+    description: The number of Triton entries restored, to pass to ./.github/actions/save-extensions.
+    value: ${{ steps.triton-count.outputs.count }}
 
 runs:
   using: composite
   steps:
     - shell: bash
-      run: echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
+      run: |
+        echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
+        echo "TRITON_CACHE_DIR=/root/.cache/triton" >> "$GITHUB_ENV"
     - id: cache
       uses: actions/cache/restore@v6
       with:
@@ -32,3 +40,14 @@ runs:
     - if: steps.cache.outputs.cache-hit == 'true'
       shell: bash
       run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} +
+    # Triton entries are content-addressed and only ever added, so any older cache is safe to start from.
+    # The run-unique key never matches: restore-keys restores the newest cache of the lane.
+    - id: triton
+      uses: actions/cache/restore@v6
+      with:
+        path: /root/.cache/triton
+        key: triton-${{ inputs.cache-key }}-${{ github.run_id }}-${{ github.run_attempt }}
+        restore-keys: triton-${{ inputs.cache-key }}-
+    - id: triton-count
+      shell: bash
+      run: echo "count=$( (find /root/.cache/triton -mindepth 1 -maxdepth 1 2>/dev/null || true) | wc -l)" >> "$GITHUB_OUTPUT"
diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-extensions/action.yml
index 09f03d45794..bd21bff8e30 100644
--- a/.github/actions/save-extensions/action.yml
+++ b/.github/actions/save-extensions/action.yml
@@ -1,11 +1,20 @@
 name: Save extensions
-description: Save the JIT-compiled extensions restored by ./.github/actions/cache-extensions. Run it with `always()` so a failing or timed-out job still
-  populates the cache (actions/cache itself only saves when the job succeeds).
+description: Save the JIT caches restored by ./.github/actions/cache-extensions. Run it with `always()` so a failing or timed-out job still populates the
+  cache (actions/cache itself only saves when the job succeeds).
 
 inputs:
   cache-key:
     description: The `cache-key` output of ./.github/actions/cache-extensions.
     required: true
+  cache-hit:
+    description: The `cache-hit` output of ./.github/actions/cache-extensions.
+    required: true
+  triton-key:
+    description: The `triton-key` output of ./.github/actions/cache-extensions.
+    required: true
+  triton-count:
+    description: The `triton-count` output of ./.github/actions/cache-extensions.
+    required: true
 
 runs:
   using: composite
@@ -15,8 +24,22 @@ runs:
     - id: built
       shell: bash
       run: find /root/.cache/torch_extensions -name '*.so' 2>/dev/null | grep -q . && echo "any=true" >> "$GITHUB_OUTPUT" || true
-    - if: steps.built.outputs.any == 'true'
+    - if: inputs.cache-hit != 'true' && steps.built.outputs.any == 'true'
       uses: actions/cache/save@v6
       with:
         path: /root/.cache/torch_extensions
         key: ${{ inputs.cache-key }}
+    # A saved Triton cache holds everything restored plus the new kernels, so save it only when
+    # a run compiled enough new ones, not on every run.
+    - id: triton
+      shell: bash
+      env:
+        RESTORED: ${{ inputs.triton-count }}
+      run: |
+        count=$( (find /root/.cache/triton -mindepth 1 -maxdepth 1 2>/dev/null || true) | wc -l)
+        if [ "$count" -ge $((${RESTORED:-0} + 20)) ]; then echo "new=true" >> "$GITHUB_OUTPUT"; fi
+    - if: steps.triton.outputs.new == 'true' && inputs.triton-key != ''
+      uses: actions/cache/save@v6
+      with:
+        path: /root/.cache/triton
+        key: ${{ inputs.triton-key }}
diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index 54f32ccce5b..fbfff80ce86 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -126,10 +126,13 @@ jobs:
           codecov-token: ${{ secrets.CODECOV_TOKEN }}
       - name: Save extensions
         # Not tied to the tests passing, see the action description
-        if: always() && steps.ext-cache.outputs.cache-hit != 'true'
+        if: always()
         uses: ./.github/actions/save-extensions
         with:
           cache-key: ${{ steps.ext-cache.outputs.cache-key }}
+          cache-hit: ${{ steps.ext-cache.outputs.cache-hit }}
+          triton-key: ${{ steps.ext-cache.outputs.triton-key }}
+          triton-count: ${{ steps.ext-cache.outputs.triton-count }}
       - name: Fail the job for an example that failed
         if: steps.example0.outputs.failed == 'true' || steps.example1.outputs.failed == 'true' || steps.example2.outputs.failed == 'true'
         run: exit 1
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 811fcc4478e..c895edc81df 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -109,10 +109,13 @@ jobs:
           python3 -m pip install nox && nox -s ${{ matrix.example }}
       - name: Save extensions
         # Not tied to the tests passing, see the action description
-        if: always() && steps.ext-cache.outputs.cache-hit != 'true'
+        if: always()
         uses: ./.github/actions/save-extensions
         with:
           cache-key: ${{ steps.ext-cache.outputs.cache-key }}
+          cache-hit: ${{ steps.ext-cache.outputs.cache-hit }}
+          triton-key: ${{ steps.ext-cache.outputs.triton-key }}
+          triton-count: ${{ steps.ext-cache.outputs.triton-count }}
       - name: Check the wheelhouse is complete
         id: wheelhouse-complete
         if: always() && matrix.example == 'gpu'

From 6e2e7b3d1773dae74a6e9363f8804d56079a4f3d Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 21 Oct 2026 00:51:25 +0000
Subject: [PATCH 2/6] Build the source-built CUDA wheels for the CI GPU's
 architecture only

mamba_ssm and causal-conv1d hardcode a -gencode flag for every architecture of
the CUDA toolkit (nine with CUDA 13) and ignore TORCH_CUDA_ARCH_LIST, so a
wheelhouse miss compiled eight architectures the CI runner never runs. The step
that already patches their setup.py to C++20 now also drops the -gencode entries
that are not in TORCH_CUDA_ARCH_LIST (the GPU jobs set it to the runner's GPU)
and fails if none matches. Without the variable nothing changes.

fast-hadamard-transform is built from a git URL, with no patch step, and still
builds every architecture.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 noxfile.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/noxfile.py b/noxfile.py
index 32ff854a0f1..1785fe8f16c 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -25,6 +25,7 @@
 
 import glob
 import os
+import re
 import shutil
 import tarfile
 import tempfile
@@ -133,6 +134,24 @@ def partial_unit(session, subset):
     "mamba_ssm": {"MAMBA_FORCE_BUILD": "TRUE"},
     "causal_conv1d": {"CAUSAL_CONV1D_FORCE_BUILD": "TRUE"},
 }
+# Their setup.py files hardcode one -gencode flag per architecture (nine with CUDA 13) and ignore TORCH_CUDA_ARCH_LIST.
+_GENCODE = re.compile(
+    r'^(?P<indent>[ \t]*)cc_flag\.append\("-gencode"\)\n[ \t]*cc_flag\.append\("arch=compute_(?P<arch>\d+),code=sm_\d+"\)',
+    re.MULTILINE,
+)
+
+
+def _only_runner_archs(setup_py):
+    """Drop the ``-gencode`` entries for the architectures not in TORCH_CUDA_ARCH_LIST, which CI sets to its GPU."""
+    archs = os.environ.get("TORCH_CUDA_ARCH_LIST")
+    if not archs:
+        return setup_py
+    keep = {re.sub(r"\D", "", arch) for arch in re.split(r"[;\s]+", archs)}
+    # ``pass`` keeps an ``if`` block valid when all of its entries are dropped
+    patched = _GENCODE.sub(lambda m: m[0] if m["arch"] in keep else f"{m['indent']}pass", setup_py)
+    if not _GENCODE.search(patched):
+        raise ValueError(f"No -gencode entry of setup.py matches TORCH_CUDA_ARCH_LIST={archs!r}")
+    return patched
 
 
 def _wheelhouse():
@@ -157,7 +176,7 @@ def _build_wheel(session, name, wheelhouse):
         source = sdist.removesuffix(".tar.gz")
         setup_py = os.path.join(source, "setup.py")
         with open(setup_py) as f:
-            patched = f.read().replace("-std=c++17", "-std=c++20")
+            patched = _only_runner_archs(f.read().replace("-std=c++17", "-std=c++20"))
         with open(setup_py, "w") as f:
             f.write(patched)
         session.run(*pip_wheel, source, env=_FORCE_BUILD[name])

From 57c4305253dffbd22e7c4f38da5e4e4f739a89a0 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 21 Oct 2026 03:19:45 +0000
Subject: [PATCH 3/6] Save the Triton cache once a week and when the kernels
 change

Saving a Triton cache whenever a run compiled 20 more entries than it restored
saved a new copy on every run in the two lanes that use torch.compile
(diffusers, speculative_decoding): each run compiles about 25 new entries that
later runs do not reuse, so a warm run saved a 24-58 MB cache again.

The key now carries the ISO week and the hash of modelopt/torch/kernels, and
the save runs only when that exact key was missing and the cache is not empty,
like the extension cache. restore-keys still restores the newest cache of the
lane, so a run starts from the last one. The cache is refreshed once a week and
when the kernels change.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .github/actions/cache-extensions/action.yml | 19 ++++++++++---------
 .github/actions/save-extensions/action.yml  | 14 ++++----------
 .github/workflows/_example_tests_runner.yml |  2 +-
 .github/workflows/gpu_tests.yml             |  2 +-
 4 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml
index e0eb96cf8a6..acf33436048 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/cache-extensions/action.yml
@@ -16,9 +16,9 @@ outputs:
   triton-key:
     description: The Triton key to pass to ./.github/actions/save-extensions.
     value: ${{ steps.triton.outputs.cache-primary-key }}
-  triton-count:
-    description: The number of Triton entries restored, to pass to ./.github/actions/save-extensions.
-    value: ${{ steps.triton-count.outputs.count }}
+  triton-hit:
+    description: Whether the Triton cache was restored from the exact key.
+    value: ${{ steps.triton.outputs.cache-hit }}
 
 runs:
   using: composite
@@ -40,14 +40,15 @@ runs:
     - if: steps.cache.outputs.cache-hit == 'true'
       shell: bash
       run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} +
-    # Triton entries are content-addressed and only ever added, so any older cache is safe to start from.
-    # The run-unique key never matches: restore-keys restores the newest cache of the lane.
+    - id: week
+      shell: bash
+      run: echo "week=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT"
+    # Triton entries are content-addressed and only ever added, so any older cache is safe to start from:
+    # restore-keys takes the newest. The week and the kernel sources are in the exact key, so a new cache
+    # is saved once a week and when the kernels change, not on every run.
     - id: triton
       uses: actions/cache/restore@v6
       with:
         path: /root/.cache/triton
-        key: triton-${{ inputs.cache-key }}-${{ github.run_id }}-${{ github.run_attempt }}
+        key: triton-${{ inputs.cache-key }}-${{ steps.week.outputs.week }}-${{ hashFiles('modelopt/torch/kernels/**/*.py') }}
         restore-keys: triton-${{ inputs.cache-key }}-
-    - id: triton-count
-      shell: bash
-      run: echo "count=$( (find /root/.cache/triton -mindepth 1 -maxdepth 1 2>/dev/null || true) | wc -l)" >> "$GITHUB_OUTPUT"
diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-extensions/action.yml
index bd21bff8e30..e2377091447 100644
--- a/.github/actions/save-extensions/action.yml
+++ b/.github/actions/save-extensions/action.yml
@@ -12,8 +12,8 @@ inputs:
   triton-key:
     description: The `triton-key` output of ./.github/actions/cache-extensions.
     required: true
-  triton-count:
-    description: The `triton-count` output of ./.github/actions/cache-extensions.
+  triton-hit:
+    description: The `triton-hit` output of ./.github/actions/cache-extensions.
     required: true
 
 runs:
@@ -29,16 +29,10 @@ runs:
       with:
         path: /root/.cache/torch_extensions
         key: ${{ inputs.cache-key }}
-    # A saved Triton cache holds everything restored plus the new kernels, so save it only when
-    # a run compiled enough new ones, not on every run.
     - id: triton
       shell: bash
-      env:
-        RESTORED: ${{ inputs.triton-count }}
-      run: |
-        count=$( (find /root/.cache/triton -mindepth 1 -maxdepth 1 2>/dev/null || true) | wc -l)
-        if [ "$count" -ge $((${RESTORED:-0} + 20)) ]; then echo "new=true" >> "$GITHUB_OUTPUT"; fi
-    - if: steps.triton.outputs.new == 'true' && inputs.triton-key != ''
+      run: test -n "$(ls -A /root/.cache/triton 2>/dev/null)" && echo "any=true" >> "$GITHUB_OUTPUT" || true
+    - if: inputs.triton-hit != 'true' && steps.triton.outputs.any == 'true' && inputs.triton-key != ''
       uses: actions/cache/save@v6
       with:
         path: /root/.cache/triton
diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index fbfff80ce86..1acac45791a 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -132,7 +132,7 @@ jobs:
           cache-key: ${{ steps.ext-cache.outputs.cache-key }}
           cache-hit: ${{ steps.ext-cache.outputs.cache-hit }}
           triton-key: ${{ steps.ext-cache.outputs.triton-key }}
-          triton-count: ${{ steps.ext-cache.outputs.triton-count }}
+          triton-hit: ${{ steps.ext-cache.outputs.triton-hit }}
       - name: Fail the job for an example that failed
         if: steps.example0.outputs.failed == 'true' || steps.example1.outputs.failed == 'true' || steps.example2.outputs.failed == 'true'
         run: exit 1
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index c895edc81df..ba44b94ddb6 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -115,7 +115,7 @@ jobs:
           cache-key: ${{ steps.ext-cache.outputs.cache-key }}
           cache-hit: ${{ steps.ext-cache.outputs.cache-hit }}
           triton-key: ${{ steps.ext-cache.outputs.triton-key }}
-          triton-count: ${{ steps.ext-cache.outputs.triton-count }}
+          triton-hit: ${{ steps.ext-cache.outputs.triton-hit }}
       - name: Check the wheelhouse is complete
         id: wheelhouse-complete
         if: always() && matrix.example == 'gpu'

From 06163d80ae7a36b8c7e916007dd25e2fb6c17393 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 21 Oct 2026 22:35:43 +0000
Subject: [PATCH 4/6] Cache the Triton autotune results too

TRITON_CACHE_DIR caches compiled kernels, but an autotuned kernel still
benchmarks its configs on first use in every process, which the warm CI runs
still paid for (the GDN fixtures). TRITON_CACHE_AUTOTUNING=1 writes the chosen
config next to the compiled kernel, inside the same directory, so a warm job
skips the benchmark as well.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .github/actions/cache-extensions/action.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml
index acf33436048..d4cc75ce313 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/cache-extensions/action.yml
@@ -27,6 +27,7 @@ runs:
       run: |
         echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
         echo "TRITON_CACHE_DIR=/root/.cache/triton" >> "$GITHUB_ENV"
+        echo "TRITON_CACHE_AUTOTUNING=1" >> "$GITHUB_ENV"
     - id: cache
       uses: actions/cache/restore@v6
       with:

From 1e9e7c8e115b90ac100eac37ab1d4d6d03157680 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Mon, 26 Oct 2026 18:43:28 +0000
Subject: [PATCH 5/6] Rename the cache actions to restore-jit-caches and
 save-jit-caches

The actions cache more than extensions now (Triton kernels), and the next cache
(Inductor, FlashInfer, numba) would make the old names worse. They are named
after the actions/cache restore and save verbs, and "JIT" is already the term
in their descriptions.

The step id is jit-caches and the steps are "Restore JIT caches" and "Save JIT
caches". The outputs and save inputs are one key/hit pair per cache
(torch-ext-key, torch-ext-hit, triton-key, triton-hit), and the restore input
cache-key is env-key, since it is the GPU and image discriminator, not a cache
key. The cache key prefixes (torch-ext-, triton-) stay, so no existing cache is
invalidated.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../action.yml                                | 32 +++++++++----------
 .../action.yml                                | 22 ++++++-------
 .github/workflows/_example_tests_runner.yml   | 19 +++++------
 .github/workflows/gpu_tests.yml               | 19 +++++------
 4 files changed, 47 insertions(+), 45 deletions(-)
 rename .github/actions/{cache-extensions => restore-jit-caches}/action.yml (64%)
 rename .github/actions/{save-extensions => save-jit-caches}/action.yml (56%)

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/restore-jit-caches/action.yml
similarity index 64%
rename from .github/actions/cache-extensions/action.yml
rename to .github/actions/restore-jit-caches/action.yml
index d4cc75ce313..c978418b605 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/restore-jit-caches/action.yml
@@ -1,20 +1,20 @@
-name: Cache extensions
-description: Restore JIT-compiled extensions (CUDA extensions, Triton kernels) cached by an earlier run. Pair with ./.github/actions/save-extensions.
+name: Restore JIT caches
+description: Restore the JIT caches (CUDA extensions, Triton kernels) saved by an earlier run. Pair with ./.github/actions/save-jit-caches.
 
 inputs:
-  cache-key:
-    description: Environment discriminator for the cache key (e.g. GPU arch + container image).
+  env-key:
+    description: Environment discriminator for the cache keys (e.g. GPU arch + lane + container image).
     required: true
 
 outputs:
-  cache-hit:
-    description: Whether the extensions were restored from the exact key.
-    value: ${{ steps.cache.outputs.cache-hit }}
-  cache-key:
-    description: The key to pass to ./.github/actions/save-extensions.
-    value: ${{ steps.cache.outputs.cache-primary-key }}
+  torch-ext-key:
+    description: The torch extensions key to pass to ./.github/actions/save-jit-caches.
+    value: ${{ steps.torch-ext.outputs.cache-primary-key }}
+  torch-ext-hit:
+    description: Whether the torch extensions cache was restored from the exact key.
+    value: ${{ steps.torch-ext.outputs.cache-hit }}
   triton-key:
-    description: The Triton key to pass to ./.github/actions/save-extensions.
+    description: The Triton key to pass to ./.github/actions/save-jit-caches.
     value: ${{ steps.triton.outputs.cache-primary-key }}
   triton-hit:
     description: Whether the Triton cache was restored from the exact key.
@@ -28,17 +28,17 @@ runs:
         echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
         echo "TRITON_CACHE_DIR=/root/.cache/triton" >> "$GITHUB_ENV"
         echo "TRITON_CACHE_AUTOTUNING=1" >> "$GITHUB_ENV"
-    - id: cache
+    - id: torch-ext
       uses: actions/cache/restore@v6
       with:
         path: /root/.cache/torch_extensions
         # Only the JIT-built sources and their build scripts, so Triton-only edits keep the cache.
-        key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**/*.cpp', 'modelopt/torch/kernels/quantization/**/*.cu',
+        key: torch-ext-${{ inputs.env-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**/*.cpp', 'modelopt/torch/kernels/quantization/**/*.cu',
           'modelopt/torch/kernels/quantization/**/*.cuh', 'modelopt/torch/kernels/quantization/**/*.h', 'modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py',
           'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py') }}
     # On a cache hit, backdate sources below the cached objects so ninja reuses them (touching
     # the objects instead would desync ninja's .ninja_deps and force a rebuild).
-    - if: steps.cache.outputs.cache-hit == 'true'
+    - if: steps.torch-ext.outputs.cache-hit == 'true'
       shell: bash
       run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} +
     - id: week
@@ -51,5 +51,5 @@ runs:
       uses: actions/cache/restore@v6
       with:
         path: /root/.cache/triton
-        key: triton-${{ inputs.cache-key }}-${{ steps.week.outputs.week }}-${{ hashFiles('modelopt/torch/kernels/**/*.py') }}
-        restore-keys: triton-${{ inputs.cache-key }}-
+        key: triton-${{ inputs.env-key }}-${{ steps.week.outputs.week }}-${{ hashFiles('modelopt/torch/kernels/**/*.py') }}
+        restore-keys: triton-${{ inputs.env-key }}-
diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-jit-caches/action.yml
similarity index 56%
rename from .github/actions/save-extensions/action.yml
rename to .github/actions/save-jit-caches/action.yml
index e2377091447..bd56a58ad6d 100644
--- a/.github/actions/save-extensions/action.yml
+++ b/.github/actions/save-jit-caches/action.yml
@@ -1,19 +1,19 @@
-name: Save extensions
-description: Save the JIT caches restored by ./.github/actions/cache-extensions. Run it with `always()` so a failing or timed-out job still populates the
-  cache (actions/cache itself only saves when the job succeeds).
+name: Save JIT caches
+description: Save the JIT caches restored by ./.github/actions/restore-jit-caches. Run it with `always()` so a failing or timed-out job still populates
+  the cache (actions/cache itself only saves when the job succeeds).
 
 inputs:
-  cache-key:
-    description: The `cache-key` output of ./.github/actions/cache-extensions.
+  torch-ext-key:
+    description: The `torch-ext-key` output of ./.github/actions/restore-jit-caches.
     required: true
-  cache-hit:
-    description: The `cache-hit` output of ./.github/actions/cache-extensions.
+  torch-ext-hit:
+    description: The `torch-ext-hit` output of ./.github/actions/restore-jit-caches.
     required: true
   triton-key:
-    description: The `triton-key` output of ./.github/actions/cache-extensions.
+    description: The `triton-key` output of ./.github/actions/restore-jit-caches.
     required: true
   triton-hit:
-    description: The `triton-hit` output of ./.github/actions/cache-extensions.
+    description: The `triton-hit` output of ./.github/actions/restore-jit-caches.
     required: true
 
 runs:
@@ -24,11 +24,11 @@ runs:
     - id: built
       shell: bash
       run: find /root/.cache/torch_extensions -name '*.so' 2>/dev/null | grep -q . && echo "any=true" >> "$GITHUB_OUTPUT" || true
-    - if: inputs.cache-hit != 'true' && steps.built.outputs.any == 'true'
+    - if: inputs.torch-ext-hit != 'true' && steps.built.outputs.any == 'true'
       uses: actions/cache/save@v6
       with:
         path: /root/.cache/torch_extensions
-        key: ${{ inputs.cache-key }}
+        key: ${{ inputs.torch-ext-key }}
     - id: triton
       shell: bash
       run: test -n "$(ls -A /root/.cache/triton 2>/dev/null)" && echo "any=true" >> "$GITHUB_OUTPUT" || true
diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index 1acac45791a..72bd4bff751 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -67,11 +67,12 @@ jobs:
         run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@eadc57871cde0c92a223d0f486ab669a2360f978 # main
-      - uses: ./.github/actions/cache-extensions
-        id: ext-cache
+      - name: Restore JIT caches
+        id: jit-caches
+        uses: ./.github/actions/restore-jit-caches
         with:
           # Per job (its examples): jobs sharing an image build different extensions, and the first to save wins the key.
-          cache-key: rtxpro6000-example-${{ steps.examples.outputs.key }}-${{ inputs.docker_image }}
+          env-key: rtxpro6000-example-${{ steps.examples.outputs.key }}-${{ inputs.docker_image }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
@@ -124,15 +125,15 @@ jobs:
           example: ${{ steps.examples.outputs.e2 }}
           allow-failure-examples: ${{ inputs.allow_failure_examples }}
           codecov-token: ${{ secrets.CODECOV_TOKEN }}
-      - name: Save extensions
+      - name: Save JIT caches
         # Not tied to the tests passing, see the action description
         if: always()
-        uses: ./.github/actions/save-extensions
+        uses: ./.github/actions/save-jit-caches
         with:
-          cache-key: ${{ steps.ext-cache.outputs.cache-key }}
-          cache-hit: ${{ steps.ext-cache.outputs.cache-hit }}
-          triton-key: ${{ steps.ext-cache.outputs.triton-key }}
-          triton-hit: ${{ steps.ext-cache.outputs.triton-hit }}
+          torch-ext-key: ${{ steps.jit-caches.outputs.torch-ext-key }}
+          torch-ext-hit: ${{ steps.jit-caches.outputs.torch-ext-hit }}
+          triton-key: ${{ steps.jit-caches.outputs.triton-key }}
+          triton-hit: ${{ steps.jit-caches.outputs.triton-hit }}
       - name: Fail the job for an example that failed
         if: steps.example0.outputs.failed == 'true' || steps.example1.outputs.failed == 'true' || steps.example2.outputs.failed == 'true'
         run: exit 1
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index ba44b94ddb6..62980cc2a4b 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -74,11 +74,12 @@ jobs:
         run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@eadc57871cde0c92a223d0f486ab669a2360f978 # main
-      - uses: ./.github/actions/cache-extensions
-        id: ext-cache
+      - name: Restore JIT caches
+        id: jit-caches
+        uses: ./.github/actions/restore-jit-caches
         with:
           # Per job: jobs sharing an image build different extensions, and the first to save wins the key.
-          cache-key: rtxpro6000-${{ matrix.example }}-${{ matrix.container_image }}
+          env-key: rtxpro6000-${{ matrix.example }}-${{ matrix.container_image }}
       - name: Restore wheels built from source
         # mamba_ssm and fast-hadamard-transform take ~10 min to compile (see noxfile.py `gpu`)
         id: wheelhouse
@@ -107,15 +108,15 @@ jobs:
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
-      - name: Save extensions
+      - name: Save JIT caches
         # Not tied to the tests passing, see the action description
         if: always()
-        uses: ./.github/actions/save-extensions
+        uses: ./.github/actions/save-jit-caches
         with:
-          cache-key: ${{ steps.ext-cache.outputs.cache-key }}
-          cache-hit: ${{ steps.ext-cache.outputs.cache-hit }}
-          triton-key: ${{ steps.ext-cache.outputs.triton-key }}
-          triton-hit: ${{ steps.ext-cache.outputs.triton-hit }}
+          torch-ext-key: ${{ steps.jit-caches.outputs.torch-ext-key }}
+          torch-ext-hit: ${{ steps.jit-caches.outputs.torch-ext-hit }}
+          triton-key: ${{ steps.jit-caches.outputs.triton-key }}
+          triton-hit: ${{ steps.jit-caches.outputs.triton-hit }}
       - name: Check the wheelhouse is complete
         id: wheelhouse-complete
         if: always() && matrix.example == 'gpu'

From 26c469ab0d60e4262a0d87bdba6c357af2a4fc4f Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Mon, 26 Oct 2026 18:43:59 +0000
Subject: [PATCH 6/6] Keep the full architecture list when TORCH_CUDA_ARCH_LIST
 has names

PyTorch accepts named lists such as Hopper (9.0+PTX), which cannot be matched
to the -gencode entries of the pinned setup.py files. The filter treated a name
as an empty architecture, removed every entry and raised. A list with any
non-numeric entry, or no list, now leaves setup.py untouched, which is the
build without the filter; numeric lists filter as before and a numeric list
that matches no entry still raises.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 noxfile.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/noxfile.py b/noxfile.py
index 1785fe8f16c..4b64827afa9 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -143,14 +143,19 @@ def partial_unit(session, subset):
 
 def _only_runner_archs(setup_py):
     """Drop the ``-gencode`` entries for the architectures not in TORCH_CUDA_ARCH_LIST, which CI sets to its GPU."""
-    archs = os.environ.get("TORCH_CUDA_ARCH_LIST")
-    if not archs:
+    archs = [
+        arch for arch in re.split(r"[;\s]+", os.environ.get("TORCH_CUDA_ARCH_LIST", "")) if arch
+    ]
+    # A named list (``Hopper``) cannot be matched to the -gencode entries: build every architecture, as without the variable
+    if not archs or not all(re.fullmatch(r"\d+\.\d+[a-z]?(\+PTX)?", arch) for arch in archs):
         return setup_py
-    keep = {re.sub(r"\D", "", arch) for arch in re.split(r"[;\s]+", archs)}
+    keep = {re.sub(r"\D", "", arch) for arch in archs}
     # ``pass`` keeps an ``if`` block valid when all of its entries are dropped
     patched = _GENCODE.sub(lambda m: m[0] if m["arch"] in keep else f"{m['indent']}pass", setup_py)
     if not _GENCODE.search(patched):
-        raise ValueError(f"No -gencode entry of setup.py matches TORCH_CUDA_ARCH_LIST={archs!r}")
+        raise ValueError(
+            f"No -gencode entry of setup.py matches TORCH_CUDA_ARCH_LIST={';'.join(archs)!r}"
+        )
     return patched