Repository navigation
Speed up CI JIT compiles: Triton cache and arch-limited CUDA wheels - #2905
Draft
shengliangxu wants to merge 4 commits into
Draft
shengliangxu wants to merge 4 commits into
shengliangxu wants to merge 4 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: true
Comment |
Collaborator
Author
|
/ok to test c038edc |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2905 +/- ##
==========================================
+ Coverage 70.16% 78.91% +8.75%
==========================================
Files 659 659
Lines 74077 74077
==========================================
+ Hits 51974 58458 +6484
+ Misses 22103 15619 -6484
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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>
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>
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>
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>
shengliangxu
force-pushed
the
shengliangx/ci-shared-cache
branch
from
October 21, 2026 22:36
c038edc to
7403c55
Compare
Collaborator
Author
|
/ok to test 7403c55 |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: CI
Three small changes that cut compile time in the GPU and example jobs. Opened as a draft to measure them in CI. The first set of CI numbers below was measured on an earlier head of this branch (c038edc, cold run then a re-run of the same sha); the current head (7403c55) adds the week-key fix and the autotune cache, validated in the last bullet.
Cache Triton kernels across runs. Every GPU and example container starts with an empty Triton cache, so each job recompiles the same kernels: the GDN fixture (about 30 compiles), the Mamba first calls, the FLA fixture, the NVFP4 FP8-sweep kernel and the first vLLM engine. #2836 added the extension cache but nothing for Triton; this is the Triton half of 7f8854a on top of its
cache-extensions/save-extensionsactions.cache-extensionsexportsTRITON_CACHE_DIR=/root/.cache/tritonand restores a Triton cache keyedtriton-<job key>-<ISO week>-<hash of modelopt/torch/kernels/**/*.py>withrestore-keyson the job key, so a run starts from the newest cache of its lane. Triton entries are content-addressed and only ever added, so an older cache is a safe base.save-extensionssaves it only when the exact key was missing and the directory is not empty, like the extension cache: once a week and when the kernels change, not on every run. An earlier version saved whenever a run compiled 20 or more new entries; the warm run showed that the two lanes usingtorch.compile(diffusers,speculative_decoding) compile about 25 run-specific entries every time and re-saved a 24-58 MB copy on every run. The extension cache keeps its exact key; its hit check moves from the workflows'ifinto the action.TRITON_CACHE_AUTOTUNING=1makes Triton store the chosen autotune config next to the compiled kernel, in the same directory, so a warm job skips the benchmark too (the GDN fixtures still paid for it).mainonce so the entries land in main's cache scope. PR and merge-queue runs read it from there (a PR scope cannot be read by other PRs, and a merge group only sees main's).Build the source-built CUDA wheels for the CI GPU's architecture only.
mamba_ssm,causal-conv1dandfast-hadamard-transformhardcode one-gencodeflag per architecture of the CUDA toolkit (nine with CUDA 13: sm_75/80/87/90/100/103/110/120/121) and ignoreTORCH_CUDA_ARCH_LIST. A wheelhouse miss in thegpujob (the first run of every PR whosenoxfile.pyhash is not in main's cache scope, and the nightly after anoxfile.pychange) therefore compiles eight architectures the RTX PRO 6000 runner never uses.The step that already patches their
setup.pyto C++20 now also drops the-gencodeentries that are not inTORCH_CUDA_ARCH_LIST(the GPU jobs set it to the runner's GPU) and raises if none matches. Without the variable nothing changes.fast-hadamard-transformis built from a git URL with no patch step and still builds every architecture. The wheels are now per architecture: the wheelhouse key (image +noxfile.pyhash) must carry the GPU type before another GPU type runs thegpujob.Testing
bash -e -o pipefailwith a missing, an empty and a populated cache directory;_only_runner_archson the realsetup.pyof the three pinned packages forTORCH_CUDA_ARCH_LIST= 12.0, 9.0, "8.0;9.0+PTX" and 12.0a (only the requested-gencodepairs remain and the result compiles; unset leaves the file untouched; an architecture without an entry raises).gpujob vs the wheelhouse-miss run of PR 2610 on the same image: the three wheel builds took 7.7 min instead of 15.4 (mamba_ssm 283 -> 74 s, causal-conv1d 436 -> 179 s, fast-hadamard unchanged at 138 s); all 1827gputests pass; the wheelhouse cache entry went from 615 MB to 121 MB.megatron_bridgeexample lane 13.3 -> 7.6 min (test_distill_vlm79 -> 20 s,test_prune_minitron_vlm74 -> 33 s,test_qad59 -> 29 s);speculative_decoding10.7 -> 7.8 min (thetorch.compilesmoke test 188 -> 96 s);gpu: the NVFP4 sweep export test 60 -> 22 s and the GDN fixture setup 65 -> 38 s;gpu_megatronabout -1 min (GDN setup 174 -> 130 s, the 75 s worker-pool starts are unchanged);gpu_vllmabout -1 min (the first engine setup, 93 -> 102 s, is not helped: vLLM's own compile cache is not cached);diffusers-0.8 min. Lane totals are noisy (llm_qatmoved -11.5 min with no cache change), so judge the per-test numbers.restore-keysand saves one new exact-key Triton cache per lane (211 MB in total; 25-70 MB for the lanes that use Triton, 0.1-0.6 MB for the lanes that barely do). A re-run ofspeculative_decodingand ofgpu_megatronrestores the exact key (hit for) and saves nothing, where the previous rule re-saved 24-58 MB on every run in thetorch.compilelanes. The autotune cache:test_gdn_qat_and_sharded_restoresetup ingpu_megatronwent from 76 s to 54 s on the re-run, with the compiled kernels already restored in both runs (thetorch.compilesmoke test inspeculative_decodingis not helped, 89 -> 92 s).Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/Asetup.pyfiles above)Additional Information
Related: #2836 (wheelhouse and extension cache), #2730 (reference measurements). Merge-queue runs only read main's cache scope; they should skip cache saves. For scale: in a survey of the GPU lanes after #2846, first-use JIT was about 160 s of a PR
gpujob (GDN Triton compile 70 s, NVFP4 sweep kernels about 47 s), the first vLLM engine boot 104 s against 24-39 s steady state, and the 4torch.compiletests ofdiffusers153 of its 209 s of pytest.