From 2cf9e946f66c6ad6110fa2217a8bfc6541370dba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 00:44:02 +0000 Subject: [PATCH 1/4] 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 --- .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 18979ce0f48..98c6f11c7f3 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -114,10 +114,13 @@ jobs: python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}" - 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: Flag allowed failure if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }} run: | diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index 0b1a801b25b..8b6c0bdf53d 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -107,10 +107,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 4cd93075de8882352862f8931e0bd52cfe83eb71 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 00:51:25 +0000 Subject: [PATCH 2/4] 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 --- 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[ \t]*)cc_flag\.append\("-gencode"\)\n[ \t]*cc_flag\.append\("arch=compute_(?P\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 80e399d2a2853ec2568eb45901245523220fa6bf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 03:19:45 +0000 Subject: [PATCH 3/4] 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 --- .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 98c6f11c7f3..a6cfeedb9b4 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -120,7 +120,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: Flag allowed failure if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }} run: | diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index 8b6c0bdf53d..cdc5dacebcc 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -113,7 +113,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 88279bdb1ba385eaa5e727b812c19c1bbd0ce3d4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 22:35:43 +0000 Subject: [PATCH 4/4] 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 --- .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: