fix(jit): invalidate stale nvcc cache artifacts by build fingerprint - #4490
zhachar14h-cell wants to merge 5 commits into
Conversation
FlashInfer's nvcc JIT cache relies on Ninja's mtime scan to decide whether a compiled module is reusable. When the same wheel is reinstalled, its source files may carry mtimes *older* than the previously compiled artifacts, so Ninja treats a stale `.so`/`.o` as up-to-date and reuses it even though the C++/CUDA bytes changed. The result is a module linked from a mix of old and new objects, which can deadlock at launch (observed on sm_121a with `fused_moe_120.so`). Add a per-module `meta.json` build fingerprint to `JitSpecNvcc` covering the FlashInfer version/git commit, wheel `RECORD` hash, Python SOABI, Torch + CUDA + libstdc++ ABI, TVM-FFI version, source-tree content hashes, the generated Ninja file, and the compile flags. `build()` now wipes the module build directory whenever the committed fingerprint is missing, corrupt, or different, then rebuilds. The fingerprint is committed atomically only after a successful compile; a failed build never marks partial artifacts valid. Ninja's fine-grained header mtime scan is preserved; this only adds a cheap content/ABI gate on top. AOT loading and `FLASHINFER_DISABLE_JIT` are unchanged. Old caches without `meta.json` get one conservative rebuild on first use after upgrade. Fixes flashinfer-ai#4489 (NVCC JIT cache can reuse stale same-version binaries after reinstall). Test plan: `pytest -q tests/jit/test_jit_cpp_ext.py` (new GPU-free tests cover source/flag/identity/ABI invalidation, corrupt/missing meta, ninja failure, AOT behavior, determinism).
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughNVCC JIT builds now use content and runtime fingerprints stored in metadata. Stale or missing metadata invalidates JIT build directories. Metadata is written only after successful compilation. Batch builds apply the same validation and persistence rules. ChangesNVCC JIT fingerprinting
Estimated code review effort: 4 (Complex) | ~45 minutes Mergeability Score: 🟡 Moderate · up to Batch JIT builds may record a valid-looking fingerprint for artifacts changed concurrently by another builder, allowing mismatched cached objects to be reused. Merge should wait for this concurrency risk to be fixed or explicitly accepted by the owner. Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant JitSpecNvcc
participant Metadata
participant Ninja
participant JITLibrary
JitSpecNvcc->>Metadata: Validate expected fingerprints
JitSpecNvcc->>Ninja: Invalidate stale directory and run build
Ninja->>JITLibrary: Produce compiled library
JitSpecNvcc->>Metadata: Atomically write metadata
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/jit/core.py`:
- Around line 212-216: Update the RECORD hashing logic in the visible
record-hash function so each flashinfer/ row includes its recorded content,
rather than hashing only the path before the first comma. Preserve filtering to
flashinfer entries, and hash the full RECORD row or at minimum the path together
with its recorded content hash so wheels with identical paths but different
bytes produce different wheel_record_hash values.
- Around line 159-191: Update the fingerprint construction around self.sources
to include a content hash for every source file, not just include_dirs. Reuse
_hash_file_sha256 or equivalent and ensure generated sources are included in the
metadata used for the JIT module identifier. Do not use _hash_source_tree_cached
or another process-lifetime memoized digest for mutable source files unless it
has content-based invalidation.
In `@tests/jit/test_jit_cpp_ext.py`:
- Around line 301-336: Update the source invalidation test using _make_nvcc_spec
to create a sentinel file in spec.build_dir after the initial build, then modify
the source and trigger the rebuild. Assert the sentinel is removed and
source_sha256 changes, rather than relying only on calls["builds"], so the test
verifies _invalidate_stale_build() removes stale build directories and fails
unless spec.sources affects the fingerprint.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ad3dc665-c1c9-4c37-ac44-4f7458579674
📒 Files selected for processing (2)
flashinfer/jit/core.pytests/jit/test_jit_cpp_ext.py
- Cover the batch precompile path (build_jit_specs) with the same fingerprint gate: wipe stale dirs before the shared ninja run and commit meta.json for every module after it. - Hash the full wheel RECORD row (path + content digest + size) so a same-version reinstall with different bytes invalidates even when paths are unchanged; use the real distribution name "apache-tvm-ffi". - Recompute the include-tree hash on every call (drop the process-wide cache) so editable-install source edits are always observed; sort os.walk dirs and files for a deterministic digest. - Add the nvcc / host-CXX compiler identity to the fingerprint (torch.version. cuda does not represent the local toolchain). - Abort instead of ignoring failures when a stale build dir cannot be wiped, and write meta.json atomically via a process-unique temp file. - Snapshot the fingerprint before building and commit that snapshot, so a mid-build source edit can never bless objects with a post-hoc fingerprint. - Make JitSpecNvcc.is_compiled require a valid meta.json for JIT artifacts (AOT artifacts remain valid by construction). - Tests: use a built.marker sentinel to assert the stale dir is really wiped; add coverage for module-source vs include-tree invalidation, the batch path, wipe failure, and is_compiled semantics.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/jit/core.py (2)
530-615: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winUpdate the JIT cache documentation.
Document the
meta.jsonfingerprint contract, stale-directory handling, and metadata commit rule inCLAUDE.mdand the relevant JIT skill documentation.As per coding guidelines: “Keep documentation synchronized with infrastructure, API, convention, error-handling, and macro changes; update
CLAUDE.mdand relevant skill documentation immediately.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/jit/core.py` around lines 530 - 615, Update CLAUDE.md and the relevant JIT skill documentation to describe the metadata contract implemented by expected_meta and _meta_matches, including the fingerprint fields, stale build-directory invalidation, treatment of directories without meta.json, and the requirement to write meta.json only after a successful build. Keep the documentation aligned with _invalidate_stale_build and the existing JIT cache behavior.Source: Coding guidelines
872-893: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftKeep each module lock until its metadata is committed.
The per-module lock ends before
run_ninja()and before_write_meta_atomic(). Another process can modify or remove the module directory after this code snapshotsexpected_meta. This batch then can commit its snapshot for an artifact produced by the other process.Acquire module locks in deterministic order and retain them through the shared Ninja run and each metadata commit. Add a multi-process regression for a batch build racing a single-spec build of the same module.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/jit/core.py` around lines 872 - 893, Update the batch-build flow around the per-module FileLock and run_ninja so every module lock is acquired in deterministic order and held until its corresponding _write_meta_atomic commit completes after the shared Ninja run. Prevent locks from being released when built_specs snapshots expected_meta; preserve the existing stale-build invalidation and metadata behavior, and add a multi-process regression covering a batch build racing a single-spec build for the same module.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/jit/core.py`:
- Around line 586-596: Update the Ninja metadata flow around expected_meta,
_meta_matches, and write_ninja so the current configuration is rendered through
one shared helper before comparison. Hash that rendered content for
ninja_content_sha256 instead of hashing the existing ninja_path, and make
write_ninja write the same rendered content; add a regression that changes a
Ninja-generation input while retaining the existing .so.
- Around line 598-615: Update _invalidate_stale_build to remove the build
directory whenever meta.json is absent or _meta_matches() reports a fingerprint
mismatch, while retaining the existing logging, failure propagation, and True
return behavior. Add a regression covering an existing .so without metadata,
verifying the directory is removed so Ninja cannot treat the artifact as up to
date.
---
Outside diff comments:
In `@flashinfer/jit/core.py`:
- Around line 530-615: Update CLAUDE.md and the relevant JIT skill documentation
to describe the metadata contract implemented by expected_meta and
_meta_matches, including the fingerprint fields, stale build-directory
invalidation, treatment of directories without meta.json, and the requirement to
write meta.json only after a successful build. Keep the documentation aligned
with _invalidate_stale_build and the existing JIT cache behavior.
- Around line 872-893: Update the batch-build flow around the per-module
FileLock and run_ninja so every module lock is acquired in deterministic order
and held until its corresponding _write_meta_atomic commit completes after the
shared Ninja run. Prevent locks from being released when built_specs snapshots
expected_meta; preserve the existing stale-build invalidation and metadata
behavior, and add a multi-process regression covering a batch build racing a
single-spec build for the same module.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 1b86d09d-0cd3-4754-b76f-2b71517443fa
📒 Files selected for processing (2)
flashinfer/jit/core.pytests/jit/test_jit_cpp_ext.py
|
Addressed the two outside-diff review findings in
Validation: 18 focused tests passed; all pre-commit hooks passed. The full test file is 36/37, with the sole failure being the pre-existing missing |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@CLAUDE.md`:
- Around line 285-299: Update the later Module Caching section’s URI description
to state that the URI identifies the module configuration, while meta.json
validates cached build artifacts against the fingerprint contract. Remove the
outdated claim that the URI directly includes source_hashes, flags, and
cuda_arch, and align the related parameter-derived wording near that section.
In `@flashinfer/jit/core.py`:
- Around line 887-888: Update the lock acquisition loop over selected_specs to
use _PersistentFileLock instead of FileLock, preserving thread_local=False so
batch module locks remain persistent across releases on the shared cache.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 26e38335-d922-427f-8528-0e9285a11d2a
📒 Files selected for processing (4)
.claude/skills/add-cuda-kernel/SKILL.mdCLAUDE.mdflashinfer/jit/core.pytests/jit/test_jit_cpp_ext.py
|
@flashinfer-bot run |
|
@yongwww The latest commit is |
|
@flashinfer-bot run |
|
CI fix pushed in |
|
Update: authorized CI was subsequently started correctly for current HEAD |
|
Hi @bkryu — would you be willing to review this PR when you have time? You authored the current |
3 participants
Summary
FlashInfer's nvcc JIT cache (under
~/.cache/flashinfer/) is validated by Ninja's file mtime scan alone. When a wheel is reinstalled at the same version, the wheel's source files may carry older mtimes than the previously compiled artifacts. Ninja then treats the stale.so/.oas up-to-date and reuses them even though the C++/CUDA source bytes changed, producing a module linked from a mix of old and new objects. On sm_121a this manifested as afused_moe_120.sothat deadlockedcutlass_fused_moeat launch (GPU idle, scheduler stuck in epoll/futex, 300s watchdog timeout).This PR adds a per-module
meta.jsonbuild fingerprint toJitSpecNvcc. The build directory is wiped and rebuilt whenever the committed fingerprint is missing, corrupt, or different; the fingerprint is committed atomically only after a successful compile, and only as a snapshot taken before the build (a mid-build source edit can never bless objects with a post-hoc fingerprint). Ninja's fine-grained header mtime scan is preserved — this only adds a cheap content/ABI gate on top. AOT loading andFLASHINFER_DISABLE_JITare unchanged. Old caches withoutmeta.jsonget one conservative rebuild on first use after upgrade.Fingerprint contents
RECORDdigest for files underflashinfer/— hashes the full row (path + contentsha256=+ size), so a same-version reinstall with different bytes invalidates even when paths are unchangedSOABIapache-tvm-ffi)FLASHINFER_NVCC/nvcc,CXX/c++, launchers)data/csrcand repo-treecsrc/for source installs), recomputed on every call so editable-install edits are always observed;os.walkdirs/files are sorted for a deterministic digestsources_sha256)Scope
trtllm_fmha_v2/generated/) are staged and restored across invalidation, while stale compiler outputs are still removed.\n* Both the single-moduleJitSpecNvcc.build()and the batch precompile pathbuild_jit_specs()run the same fingerprint gate: stale dirs are wiped before building andmeta.jsonis committed after a successful shared ninja run, so a batch/AOT build can never copy stale artifacts. Module locks are acquired in deterministic order and held through the shared Ninja run and every metadata commit.meta.jsonis written via a process-unique temp file +os.replace.JitSpecNvcc.is_compilednow requires a validmeta.jsonfor JIT artifacts (AOT artifacts remain valid by construction).Test plan
pytest -q tests/jit/test_jit_cpp_ext.py— new GPU-free tests cover: source content change with unchanged (older) mtime, module-source vs include-tree invalidation independently, compile-flag change, missing/corruptmeta.json, matching fingerprint preserving the build dir (Ninja owns incrementality), generated-source preservation during invalidation, failed Ninja not committing metadata, AOT path ignoring meta, the batch precompile path, current-rendered-Ninja invalidation, cross-process batch/single lock exclusion, wipe-failure abort,is_compiledsemantics, JSON serializability, deterministic metadata. Abuilt.markersentinel asserts the stale dir is actually wiped.flashinfer/data/csrc/batch_prefill_customize_config.jinja.meta.jsonis conservatively rebuilt once and then stable; the batch path behaves identically.Fixes #4489
Notes
MAX_JOBS=1.Summary by CodeRabbit
Summary by CodeRabbit
Bug Fixes
Documentation