diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml deleted file mode 100644 index a25dcb5fafc..00000000000 --- a/.github/actions/cache-extensions/action.yml +++ /dev/null @@ -1,34 +0,0 @@ -name: Cache extensions -description: Restore JIT-compiled extensions (e.g. CUDA extensions) cached by an earlier run. Pair with ./.github/actions/save-extensions. - -inputs: - cache-key: - description: Environment discriminator for the cache key (e.g. GPU arch + 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 }} - -runs: - using: composite - steps: - - shell: bash - run: echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV" - - id: cache - 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', - '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' - shell: bash - run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} + diff --git a/.github/actions/restore-jit-caches/action.yml b/.github/actions/restore-jit-caches/action.yml new file mode 100644 index 00000000000..c978418b605 --- /dev/null +++ b/.github/actions/restore-jit-caches/action.yml @@ -0,0 +1,55 @@ +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: + env-key: + description: Environment discriminator for the cache keys (e.g. GPU arch + lane + container image). + required: true + +outputs: + 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-jit-caches. + 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" + echo "TRITON_CACHE_DIR=/root/.cache/triton" >> "$GITHUB_ENV" + echo "TRITON_CACHE_AUTOTUNING=1" >> "$GITHUB_ENV" + - 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.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.torch-ext.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.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-extensions/action.yml deleted file mode 100644 index 09f03d45794..00000000000 --- a/.github/actions/save-extensions/action.yml +++ /dev/null @@ -1,22 +0,0 @@ -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). - -inputs: - cache-key: - description: The `cache-key` output of ./.github/actions/cache-extensions. - required: true - -runs: - using: composite - steps: - # A job cancelled before building anything must not claim the key with an empty cache: - # later runs would restore it as an exact hit and never save again. - - 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' - uses: actions/cache/save@v6 - with: - path: /root/.cache/torch_extensions - key: ${{ inputs.cache-key }} diff --git a/.github/actions/save-jit-caches/action.yml b/.github/actions/save-jit-caches/action.yml new file mode 100644 index 00000000000..bd56a58ad6d --- /dev/null +++ b/.github/actions/save-jit-caches/action.yml @@ -0,0 +1,39 @@ +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: + torch-ext-key: + description: The `torch-ext-key` output of ./.github/actions/restore-jit-caches. + required: true + 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/restore-jit-caches. + required: true + triton-hit: + description: The `triton-hit` output of ./.github/actions/restore-jit-caches. + required: true + +runs: + using: composite + steps: + # A job cancelled before building anything must not claim the key with an empty cache: + # later runs would restore it as an exact hit and never save again. + - 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.torch-ext-hit != 'true' && steps.built.outputs.any == 'true' + uses: actions/cache/save@v6 + with: + path: /root/.cache/torch_extensions + 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 + - 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 54f32ccce5b..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,12 +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() && steps.ext-cache.outputs.cache-hit != 'true' - uses: ./.github/actions/save-extensions + if: always() + uses: ./.github/actions/save-jit-caches with: - cache-key: ${{ steps.ext-cache.outputs.cache-key }} + 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 811fcc4478e..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,12 +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() && steps.ext-cache.outputs.cache-hit != 'true' - uses: ./.github/actions/save-extensions + if: always() + uses: ./.github/actions/save-jit-caches with: - cache-key: ${{ steps.ext-cache.outputs.cache-key }} + 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' diff --git a/noxfile.py b/noxfile.py index 32ff854a0f1..4b64827afa9 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,29 @@ 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 = [ + 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 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={';'.join(archs)!r}" + ) + return patched def _wheelhouse(): @@ -157,7 +181,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])