diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml index a25dcb5fafc..d4cc75ce313 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,21 @@ 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-hit: + description: Whether the Triton cache was restored from the exact key. + value: ${{ steps.triton.outputs.cache-hit }} 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" + echo "TRITON_CACHE_AUTOTUNING=1" >> "$GITHUB_ENV" - id: cache uses: actions/cache/restore@v6 with: @@ -32,3 +41,15 @@ runs: - if: steps.cache.outputs.cache-hit == 'true' shell: bash run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} + + - 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 }}-${{ steps.week.outputs.week }}-${{ hashFiles('modelopt/torch/kernels/**/*.py') }} + restore-keys: triton-${{ inputs.cache-key }}- diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-extensions/action.yml index 09f03d45794..e2377091447 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-hit: + description: The `triton-hit` output of ./.github/actions/cache-extensions. + required: true runs: using: composite @@ -15,8 +24,16 @@ 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 }} + - id: triton + shell: bash + 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 + key: ${{ inputs.triton-key }} diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index 18979ce0f48..a6cfeedb9b4 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-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 0b1a801b25b..cdc5dacebcc 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-hit: ${{ steps.ext-cache.outputs.triton-hit }} - name: Check the wheelhouse is complete id: wheelhouse-complete if: always() && matrix.example == 'gpu' 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])