fix: avoid re-downloading BMM export headers when flashinfer-cubin is installed - #2903
Conversation
|
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:
📝 WalkthroughWalkthroughRemoved the per-file SHA256-validated download/cache helper and replaced bulk header-download flow with per-header Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@flashinfer/jit/cubin_loader.py`:
- Around line 217-225: The fast-path copies the cached artifact with
shutil.copy(local_artifact_path, file_path), bypassing the module’s
lock/temp-file/atomic-replace behavior and risking partial writes; replace that
copy with the same atomic-write flow used by download_file (or a helper that
acquires the module lock, writes to a temp file, verifies sha256, then
os.replace into file_path) so load_cubin(local_artifact_path, sha256) remains
used to validate the artifact but the final placement of file_path is done
atomically (avoid calling shutil.copy directly).
🪄 Autofix (Beta)
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
Run ID: f164904c-e71e-405e-b8c2-f5e463b9a3e6
📥 Commits
Reviewing files that changed from the base of the PR and between 31b63bc and 1ee8dde48f23157eedf609ec6b23c83c9cf6830a.
📒 Files selected for processing (1)
flashinfer/jit/cubin_loader.py
1ee8dde to
d791004
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the cubin loader to attempt to copy files from a local artifact directory before falling back to a network download. The review identifies a race condition in the implementation of the local file copy and suggests using the existing download_file utility to ensure atomic, concurrency-safe file operations.
I am having trouble creating individual review comments. Click here to see my feedback.
flashinfer/jit/cubin_loader.py (224-225)
The use of shutil.copy here is not atomic and not safe for concurrent execution. If multiple processes execute this code simultaneously for the same file_path, it can lead to a race condition where one process might read a partially written file from another process.
To ensure atomicity and prevent race conditions, you should use the existing download_file utility. It is designed for concurrency-safe file operations by using a file lock and an atomic rename operation. This change will also make the behavior consistent with how network files are handled. Using download_file also provides a graceful fallback to network download if the local copy fails for some reason.
if download_file(local_artifact_path, file_path, session=session):
return local_file
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@flashinfer/jit/cubin_loader.py`:
- Around line 278-284: The loop that calls get_cubin(uri_path, file_hash) for
each file (iterating header_files and using header_path) can attempt to write
nested destination files into a non-existent directory, causing
FileNotFoundError during the network fallback; ensure the canonical destination
parent directories are created before the first fallback by creating the
destination directory (derived from header_path and file name) with
os.makedirs(..., exist_ok=True) or equivalent prior to calling
get_cubin()/download_file() inside the for loop so get_cubin() or
download_file() can write the file successfully.
- Around line 286-300: The symlink removal and creation around header_dest_dir
is racy; instead, serialize replacement by creating the new symlink atomically
and then swapping it into place: ensure header_dest_dir.parent exists, create a
temporary symlink (unique name) that points to artifact_dir, then atomically
replace the target with os.replace/Path.replace from the temp symlink to
header_dest_dir so concurrent workers cannot observe a missing or
partially-updated entry; reference header_dest_dir, artifact_dir,
FLASHINFER_CUBIN_DIR and header_path and ensure you remove the temp symlink on
failure and handle existing correct symlink early-return as currently
implemented.
🪄 Autofix (Beta)
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
Run ID: ee6b3161-10bc-40a6-a156-fa0cdedce788
📥 Commits
Reviewing files that changed from the base of the PR and between 1ee8dde48f23157eedf609ec6b23c83c9cf6830a and d791004e52866d2e57b0dabda9e74096e943da74.
📒 Files selected for processing (1)
flashinfer/jit/cubin_loader.py
| # Symlink header_dest_dir -> artifact directory so C++ includes resolve. | ||
| artifact_dir = FLASHINFER_CUBIN_DIR / header_path | ||
| if header_dest_dir.is_symlink() or header_dest_dir.exists(): | ||
| if ( | ||
| header_dest_dir.is_symlink() | ||
| and header_dest_dir.resolve() == artifact_dir.resolve() | ||
| ): | ||
| return # already correct | ||
| # Stale symlink or directory from a previous version; remove it. | ||
| if header_dest_dir.is_symlink() or header_dest_dir.is_file(): | ||
| header_dest_dir.unlink() | ||
| else: | ||
| shutil.rmtree(header_dest_dir) | ||
| header_dest_dir.parent.mkdir(parents=True, exist_ok=True) | ||
| header_dest_dir.symlink_to(artifact_dir) |
There was a problem hiding this comment.
Serialize the symlink replacement.
Lines 288-300 do an unlocked unlink/rmtree + symlink_to sequence. If two TP workers hit this path at startup, one can delete the entry after the other has checked it, or both can race on symlink_to(), which can abort JIT setup with FileExistsError / FileNotFoundError.
🔒 Proposed fix
# Symlink header_dest_dir -> artifact directory so C++ includes resolve.
artifact_dir = FLASHINFER_CUBIN_DIR / header_path
- if header_dest_dir.is_symlink() or header_dest_dir.exists():
- if (
- header_dest_dir.is_symlink()
- and header_dest_dir.resolve() == artifact_dir.resolve()
- ):
- return # already correct
- # Stale symlink or directory from a previous version; remove it.
- if header_dest_dir.is_symlink() or header_dest_dir.is_file():
- header_dest_dir.unlink()
- else:
- shutil.rmtree(header_dest_dir)
header_dest_dir.parent.mkdir(parents=True, exist_ok=True)
- header_dest_dir.symlink_to(artifact_dir)
+ with filelock.FileLock(f"{header_dest_dir}.lock", timeout=30):
+ if (
+ header_dest_dir.is_symlink()
+ and header_dest_dir.resolve() == artifact_dir.resolve()
+ ):
+ return # already correct
+ # Stale symlink or directory from a previous version; remove it.
+ if header_dest_dir.is_symlink() or header_dest_dir.is_file():
+ header_dest_dir.unlink()
+ elif header_dest_dir.exists():
+ shutil.rmtree(header_dest_dir)
+ header_dest_dir.symlink_to(artifact_dir)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@flashinfer/jit/cubin_loader.py` around lines 286 - 300, The symlink removal
and creation around header_dest_dir is racy; instead, serialize replacement by
creating the new symlink atomically and then swapping it into place: ensure
header_dest_dir.parent exists, create a temporary symlink (unique name) that
points to artifact_dir, then atomically replace the target with
os.replace/Path.replace from the temp symlink to header_dest_dir so concurrent
workers cannot observe a missing or partially-updated entry; reference
header_dest_dir, artifact_dir, FLASHINFER_CUBIN_DIR and header_path and ensure
you remove the temp symlink on failure and handle existing correct symlink
early-return as currently implemented.
ed84e32 to
9fff8b3
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/test_download_trtllm_headers.py`:
- Around line 66-155: Add a new test that covers the cache-miss/download
fallback: simulate FLASHINFER_CUBIN_DIR without the header_path populated (do
not call _populate_artifact_dir), monkeypatch
flashinfer.jit.cubin_loader.download_file to a stub that writes expected header
files into the artifact path (or a tmp dir) and records calls, then call
download_trtllm_headers("bmm", header_dest_dir, header_path, artifact_path,
checksums_txt) and assert download_file was invoked for the missing files, the
header_dest_dir becomes a symlink, and the symlinked headers contain the
expected contents; reference functions/vars download_trtllm_headers,
download_file, FLASHINFER_CUBIN_DIR, and header_dest_dir in the test.
🪄 Autofix (Beta)
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
Run ID: 33b87325-e855-432c-8a26-0a68887bb96b
📥 Commits
Reviewing files that changed from the base of the PR and between d791004e52866d2e57b0dabda9e74096e943da74 and 9fff8b386c6213ed3df3499ca5528244b089f908.
📒 Files selected for processing (1)
tests/test_download_trtllm_headers.py
| def test_no_download_when_files_present(monkeypatch, tmp_path): | ||
| """When files exist at the artifact path, download_file must not be called.""" | ||
| artifact_path = "deadbeef/batched_gemm-abc-def/" | ||
| header_path = f"{artifact_path}include/trtllmGen_bmm_export" | ||
| cubin_dir = tmp_path / "cubins" | ||
| cubin_dir.mkdir() | ||
|
|
||
| file_contents, checksums_txt = _build_checksums_and_files(BMM_HEADER_FILES) | ||
| _populate_artifact_dir(cubin_dir, header_path, BMM_HEADER_FILES, file_contents) | ||
|
|
||
| import flashinfer.jit.cubin_loader as cl | ||
|
|
||
| monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir) | ||
|
|
||
| header_dest_dir = ( | ||
| cubin_dir / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export" | ||
| ) | ||
|
|
||
| with patch.object(cl, "download_file", _download_file_should_not_be_called): | ||
| download_trtllm_headers( | ||
| "bmm", header_dest_dir, header_path, artifact_path, checksums_txt | ||
| ) | ||
|
|
||
| # Verify the symlink was created and points to the right place | ||
| assert header_dest_dir.is_symlink() | ||
| assert header_dest_dir.resolve() == (cubin_dir / header_path).resolve() | ||
|
|
||
| # Verify all headers are accessible through the symlink | ||
| for f in BMM_HEADER_FILES: | ||
| assert (header_dest_dir / f).exists(), f"Header {f} not accessible via symlink" | ||
| assert (header_dest_dir / f).read_bytes() == file_contents[f] | ||
|
|
||
|
|
||
| def test_stale_directory_replaced_by_symlink(monkeypatch, tmp_path): | ||
| """A pre-existing real directory at header_dest_dir gets replaced by a symlink.""" | ||
| artifact_path = "deadbeef/batched_gemm-abc-def/" | ||
| header_path = f"{artifact_path}include/trtllmGen_bmm_export" | ||
| cubin_dir = tmp_path / "cubins" | ||
| cubin_dir.mkdir() | ||
|
|
||
| file_contents, checksums_txt = _build_checksums_and_files(BMM_HEADER_FILES) | ||
| _populate_artifact_dir(cubin_dir, header_path, BMM_HEADER_FILES, file_contents) | ||
|
|
||
| import flashinfer.jit.cubin_loader as cl | ||
|
|
||
| monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir) | ||
|
|
||
| header_dest_dir = ( | ||
| cubin_dir / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export" | ||
| ) | ||
| # Simulate leftover directory from the old download_trtllm_headers behavior | ||
| header_dest_dir.mkdir(parents=True, exist_ok=True) | ||
| (header_dest_dir / "stale_file.h").write_text("old") | ||
|
|
||
| with patch.object(cl, "download_file", _download_file_should_not_be_called): | ||
| download_trtllm_headers( | ||
| "bmm", header_dest_dir, header_path, artifact_path, checksums_txt | ||
| ) | ||
|
|
||
| assert header_dest_dir.is_symlink() | ||
| assert not (header_dest_dir / "stale_file.h").exists() | ||
|
|
||
|
|
||
| def test_idempotent(monkeypatch, tmp_path): | ||
| """Calling download_trtllm_headers twice is a no-op the second time.""" | ||
| artifact_path = "deadbeef/batched_gemm-abc-def/" | ||
| header_path = f"{artifact_path}include/trtllmGen_bmm_export" | ||
| cubin_dir = tmp_path / "cubins" | ||
| cubin_dir.mkdir() | ||
|
|
||
| file_contents, checksums_txt = _build_checksums_and_files(BMM_HEADER_FILES) | ||
| _populate_artifact_dir(cubin_dir, header_path, BMM_HEADER_FILES, file_contents) | ||
|
|
||
| import flashinfer.jit.cubin_loader as cl | ||
|
|
||
| monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir) | ||
|
|
||
| header_dest_dir = ( | ||
| cubin_dir / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export" | ||
| ) | ||
|
|
||
| with patch.object(cl, "download_file", _download_file_should_not_be_called): | ||
| download_trtllm_headers( | ||
| "bmm", header_dest_dir, header_path, artifact_path, checksums_txt | ||
| ) | ||
| download_trtllm_headers( | ||
| "bmm", header_dest_dir, header_path, artifact_path, checksums_txt | ||
| ) | ||
|
|
||
| assert header_dest_dir.is_symlink() |
There was a problem hiding this comment.
Add explicit coverage for the cache-miss/download fallback path.
Current tests only validate the “all files already local” path (download_file is always forced to fail). That leaves the missing-file flow untested in this file.
✅ Suggested test addition
+import pathlib
@@
+def test_download_when_files_missing(monkeypatch, tmp_path):
+ """When artifact files are missing, download_file should be used to populate them."""
+ artifact_path = "deadbeef/batched_gemm-abc-def/"
+ header_path = f"{artifact_path}include/trtllmGen_bmm_export"
+ cubin_dir = tmp_path / "cubins"
+ cubin_dir.mkdir()
+
+ file_contents, checksums_txt = _build_checksums_and_files(BMM_HEADER_FILES)
+
+ import flashinfer.jit.cubin_loader as cl
+ monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir)
+
+ def _fake_download_file(uri, dest, session=None):
+ rel = pathlib.Path(dest).relative_to(cubin_dir).as_posix()
+ name = rel.removeprefix(f"{header_path}/")
+ p = pathlib.Path(dest)
+ p.parent.mkdir(parents=True, exist_ok=True)
+ p.write_bytes(file_contents[name])
+
+ header_dest_dir = (
+ cubin_dir / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export"
+ )
+
+ with patch.object(cl, "download_file", side_effect=_fake_download_file) as mocked:
+ download_trtllm_headers(
+ "bmm", header_dest_dir, header_path, artifact_path, checksums_txt
+ )
+
+ assert mocked.call_count > 0
+ assert header_dest_dir.is_symlink()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_download_trtllm_headers.py` around lines 66 - 155, Add a new test
that covers the cache-miss/download fallback: simulate FLASHINFER_CUBIN_DIR
without the header_path populated (do not call _populate_artifact_dir),
monkeypatch flashinfer.jit.cubin_loader.download_file to a stub that writes
expected header files into the artifact path (or a tmp dir) and records calls,
then call download_trtllm_headers("bmm", header_dest_dir, header_path,
artifact_path, checksums_txt) and assert download_file was invoked for the
missing files, the header_dest_dir becomes a symlink, and the symlinked headers
contain the expected contents; reference functions/vars download_trtllm_headers,
download_file, FLASHINFER_CUBIN_DIR, and header_dest_dir in the test.
e6673bf to
60961d8
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/test_cubin_local_cache.py (1)
37-39: Consider extracting repeated setup to a pytest fixture.The pattern of importing
flashinfer.jit.cubin_loader as cland monkeypatchingFLASHINFER_CUBIN_DIRis repeated in every test. A fixture could reduce this boilerplate:♻️ Optional refactor using a fixture
import pytest `@pytest.fixture` def cubin_loader_with_tmp_dir(monkeypatch, tmp_path): """Provides cubin_loader module with FLASHINFER_CUBIN_DIR set to a temp directory.""" import flashinfer.jit.cubin_loader as cl cubin_dir = tmp_path / "cubins" cubin_dir.mkdir() monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir) return cl, cubin_dir def test_get_cubin_no_download_when_file_present(cubin_loader_with_tmp_dir): """get_cubin() returns local bytes without calling download_file.""" cl, cubin_dir = cubin_loader_with_tmp_dir # ... rest of test using cl and cubin_dirThis is purely optional—the current approach is clear and explicit about what each test sets up.
Also applies to: 60-62, 83-85, 106-108, 185-187, 218-220, 248-250
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_cubin_local_cache.py` around lines 37 - 39, Extract the repeated import and monkeypatch into a pytest fixture that imports flashinfer.jit.cubin_loader as cl, creates a temporary cubin_dir (using tmp_path), monkeypatches cl.FLASHINFER_CUBIN_DIR to that dir, and returns (cl, cubin_dir); then update tests that currently import cl and call monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir) to accept the fixture (e.g., cubin_loader_with_tmp_dir) and use the returned cl and cubin_dir instead.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/test_cubin_local_cache.py`:
- Around line 37-39: Extract the repeated import and monkeypatch into a pytest
fixture that imports flashinfer.jit.cubin_loader as cl, creates a temporary
cubin_dir (using tmp_path), monkeypatches cl.FLASHINFER_CUBIN_DIR to that dir,
and returns (cl, cubin_dir); then update tests that currently import cl and call
monkeypatch.setattr(cl, "FLASHINFER_CUBIN_DIR", cubin_dir) to accept the fixture
(e.g., cubin_loader_with_tmp_dir) and use the returned cl and cubin_dir instead.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b9eb6734-8b52-4a80-933a-b71e7ce0310d
📥 Commits
Reviewing files that changed from the base of the PR and between 9fff8b386c6213ed3df3499ca5528244b089f908 and e6673bf4b79aa3f9725dca3277ad69d056c2c9da.
📒 Files selected for processing (1)
tests/test_cubin_local_cache.py
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
flashinfer/jit/cubin_loader.py (1)
222-242:⚠️ Potential issue | 🟠 MajorRace condition in symlink creation/replacement remains unaddressed.
The
ensure_symlink()function performs an unlocked sequence ofunlink()/rmtree()followed bysymlink_to(). When multiple TP workers start concurrently, they can race on this sequence, causingFileExistsErrororFileNotFoundErrorif one worker deletes while another is creating.Consider using atomic symlink replacement or a file lock:
🔒 Proposed fix using atomic symlink replacement
def ensure_symlink( link: Union[str, pathlib.Path], target: Union[str, pathlib.Path] ) -> None: """Create or update a symlink, removing any stale file/directory at *link*. This is used to map C++ include paths (e.g. ``CUBIN_DIR/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export``) to the canonical artifact directory where ``get_cubin()`` stores downloaded files. """ link = pathlib.Path(link) target = pathlib.Path(target) + link.parent.mkdir(parents=True, exist_ok=True) + + # Fast path: already correct if link.is_symlink() or link.exists(): if link.is_symlink() and link.resolve() == target.resolve(): return # already correct - # Stale symlink or directory from a previous version; remove it. - if link.is_symlink() or link.is_file(): - link.unlink() - else: - shutil.rmtree(link) - link.parent.mkdir(parents=True, exist_ok=True) - link.symlink_to(target) + + # Atomic replacement via temp symlink + os.replace + import uuid + tmp_link = link.parent / f".{link.name}.{uuid.uuid4().hex}.tmp" + try: + tmp_link.symlink_to(target) + # os.replace is atomic on POSIX; on Windows it may fail if link is a directory + if link.is_dir() and not link.is_symlink(): + shutil.rmtree(link) + os.replace(tmp_link, link) + except Exception: + if tmp_link.exists() or tmp_link.is_symlink(): + tmp_link.unlink() + raise🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/jit/cubin_loader.py` around lines 222 - 242, The ensure_symlink function has a race between removing a stale path and creating the symlink; fix it by making symlink creation atomic: create a temporary symlink (e.g., in the same directory with a unique name) that points to the target and then atomically replace the final link with os.replace (or pathlib.Path.replace) so other processes either see the old or new symlink but never a missing/broken state; alternatively, guard the remove+create sequence with a filesystem lock on link.parent (using a lockfile) to serialize concurrent workers—apply this change inside ensure_symlink to replace the unlink()/rmtree() + symlink_to() sequence.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@flashinfer/jit/fused_moe.py`:
- Around line 256-268: The loop that calls get_cubin(...) for headers under
bmm_export_path can fail because download_file() doesn't create parent
directories; before each get_cubin(...) call (inside the loop over
BMM_EXPORT_HEADERS in fused_moe.py) ensure the target parent directory exists
(mkdir(parents=True, exist_ok=True)) for the path where the header will be
written (derive it from include_path / bmm_export_path / header) so nested paths
like "trtllm/gen/CommonUtils.h" won't raise FileNotFoundError; update the code
around the bmm_export_path loop (referencing BMM_EXPORT_HEADERS, get_cubin(),
get_meta_hash(checksum, header), and ensure_symlink) to create the parent
directories prior to downloading each header.
In `@flashinfer/jit/moe_utils.py`:
- Around line 44-58: Duplicate header-fetching logic in moe_utils.py and
fused_moe.py should be consolidated: create a helper
download_bmm_export_headers(...) in cubin_loader.py that accepts artifact_path,
checksum_hash, and bmm_export_headers and performs the checksum retrieval (using
get_cubin and get_meta_hash), loops over BMM_EXPORT_HEADERS, creates parent
directories (header_full_path.parent.mkdir(parents=True, exist_ok=True)) before
calling get_cubin for each header to handle nested paths, and then calls
ensure_symlink to link into
jit_env.FLASHINFER_CUBIN_DIR/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export;
replace the duplicated blocks in moe_utils.py and fused_moe.py with a call to
download_bmm_export_headers(ArtifactPath.TRTLLM_GEN_BMM,
CheckSumHash.TRTLLM_GEN_BMM, BMM_EXPORT_HEADERS).
---
Duplicate comments:
In `@flashinfer/jit/cubin_loader.py`:
- Around line 222-242: The ensure_symlink function has a race between removing a
stale path and creating the symlink; fix it by making symlink creation atomic:
create a temporary symlink (e.g., in the same directory with a unique name) that
points to the target and then atomically replace the final link with os.replace
(or pathlib.Path.replace) so other processes either see the old or new symlink
but never a missing/broken state; alternatively, guard the remove+create
sequence with a filesystem lock on link.parent (using a lockfile) to serialize
concurrent workers—apply this change inside ensure_symlink to replace the
unlink()/rmtree() + symlink_to() sequence.
🪄 Autofix (Beta)
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
Run ID: 5bb200d6-22b4-44bc-a987-ff871fa25a57
📥 Commits
Reviewing files that changed from the base of the PR and between e6673bf4b79aa3f9725dca3277ad69d056c2c9da and 60961d8a6e15fa58af40adf498b2e17c55b0be5e.
📒 Files selected for processing (4)
flashinfer/jit/cubin_loader.pyflashinfer/jit/fused_moe.pyflashinfer/jit/moe_utils.pytests/test_cubin_local_cache.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_cubin_local_cache.py
| # Fetch BMM export headers via get_cubin() and symlink for C++ includes. | ||
| bmm_export_path = f"{include_path}/trtllmGen_bmm_export" | ||
| for header in BMM_EXPORT_HEADERS: | ||
| h = get_cubin(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header)) | ||
| assert h, f"{header} not found" | ||
| ensure_symlink( | ||
| jit_env.FLASHINFER_CUBIN_DIR | ||
| / "flashinfer" | ||
| / "trtllm" | ||
| / "batched_gemm" | ||
| / "trtllmGen_bmm_export" | ||
| ) | ||
|
|
||
| download_trtllm_headers( | ||
| "bmm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_BMM, checksum | ||
| / "trtllmGen_bmm_export", | ||
| jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path, | ||
| ) |
There was a problem hiding this comment.
Create parent directories before downloading nested header files.
The loop calls get_cubin() for headers with nested paths (e.g., trtllm/gen/CommonUtils.h). The download_file() function does not create parent directories, so on a clean cache the first download for a nested path will fail with FileNotFoundError.
🐛 Proposed fix to ensure directories exist
# Fetch BMM export headers via get_cubin() and symlink for C++ includes.
bmm_export_path = f"{include_path}/trtllmGen_bmm_export"
for header in BMM_EXPORT_HEADERS:
+ header_full_path = jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path / header
+ header_full_path.parent.mkdir(parents=True, exist_ok=True)
h = get_cubin(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header))
assert h, f"{header} not found"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@flashinfer/jit/fused_moe.py` around lines 256 - 268, The loop that calls
get_cubin(...) for headers under bmm_export_path can fail because
download_file() doesn't create parent directories; before each get_cubin(...)
call (inside the loop over BMM_EXPORT_HEADERS in fused_moe.py) ensure the target
parent directory exists (mkdir(parents=True, exist_ok=True)) for the path where
the header will be written (derive it from include_path / bmm_export_path /
header) so nested paths like "trtllm/gen/CommonUtils.h" won't raise
FileNotFoundError; update the code around the bmm_export_path loop (referencing
BMM_EXPORT_HEADERS, get_cubin(), get_meta_hash(checksum, header), and
ensure_symlink) to create the parent directories prior to downloading each
header.
| checksum = get_cubin( | ||
| f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt", CheckSumHash.TRTLLM_GEN_BMM | ||
| ) | ||
| bmm_export_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export" | ||
| for header in BMM_EXPORT_HEADERS: | ||
| h = get_cubin(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header)) | ||
| assert h, f"{header} not found" | ||
| ensure_symlink( | ||
| jit_env.FLASHINFER_CUBIN_DIR | ||
| / "flashinfer" | ||
| / "trtllm" | ||
| / "batched_gemm" | ||
| / "trtllmGen_bmm_export", | ||
| f"{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export", | ||
| ArtifactPath.TRTLLM_GEN_BMM, | ||
| get_cubin( | ||
| f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt", CheckSumHash.TRTLLM_GEN_BMM | ||
| ), | ||
| jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path, | ||
| ) |
There was a problem hiding this comment.
Duplicate header-fetching logic with fused_moe.py; also missing parent directory creation.
This block is nearly identical to fused_moe.py lines 243-268. Consider extracting into a shared helper (e.g., download_bmm_export_headers() in cubin_loader.py) to avoid duplication and ensure consistent behavior.
Additionally, the same issue applies here: nested header paths like trtllm/gen/CommonUtils.h require parent directory creation before get_cubin() is called.
,
♻️ Proposed refactor: extract shared helper
In cubin_loader.py, add a helper:
def download_bmm_export_headers(
artifact_path: str,
checksum_hash: str,
bmm_export_headers: list,
) -> None:
"""Download BMM export headers and create symlink for C++ includes."""
from .env import FLASHINFER_CUBIN_DIR
checksum = get_cubin(f"{artifact_path}/checksums.txt", checksum_hash)
assert checksum, f"Failed to get checksums.txt from {artifact_path}"
bmm_export_path = f"{artifact_path}/include/trtllmGen_bmm_export"
for header in bmm_export_headers:
header_full_path = FLASHINFER_CUBIN_DIR / bmm_export_path / header
header_full_path.parent.mkdir(parents=True, exist_ok=True)
h = get_cubin(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header))
assert h, f"{header} not found"
ensure_symlink(
FLASHINFER_CUBIN_DIR / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export",
FLASHINFER_CUBIN_DIR / bmm_export_path,
)Then in both fused_moe.py and moe_utils.py:
from .cubin_loader import download_bmm_export_headers
from .fused_moe import BMM_EXPORT_HEADERS
from ..artifacts import ArtifactPath, CheckSumHash
download_bmm_export_headers(ArtifactPath.TRTLLM_GEN_BMM, CheckSumHash.TRTLLM_GEN_BMM, BMM_EXPORT_HEADERS)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@flashinfer/jit/moe_utils.py` around lines 44 - 58, Duplicate header-fetching
logic in moe_utils.py and fused_moe.py should be consolidated: create a helper
download_bmm_export_headers(...) in cubin_loader.py that accepts artifact_path,
checksum_hash, and bmm_export_headers and performs the checksum retrieval (using
get_cubin and get_meta_hash), loops over BMM_EXPORT_HEADERS, creates parent
directories (header_full_path.parent.mkdir(parents=True, exist_ok=True)) before
calling get_cubin for each header to handle nested paths, and then calls
ensure_symlink to link into
jit_env.FLASHINFER_CUBIN_DIR/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export;
replace the duplicated blocks in moe_utils.py and fused_moe.py with a call to
download_bmm_export_headers(ArtifactPath.TRTLLM_GEN_BMM,
CheckSumHash.TRTLLM_GEN_BMM, BMM_EXPORT_HEADERS).
60961d8 to
74c8305
Compare
… installed `download_trtllm_headers()` stored BMM export headers at a custom path that didn't match where `download_artifacts()` (and the flashinfer-cubin wheel) puts them. This caused ~17 header files to be re-downloaded from the network on every startup. Replace `download_trtllm_headers()` and `get_file()` with two unified primitives: - `get_artifact()` (renamed from `get_cubin()`) fetches any artifact into the canonical path, reusing flashinfer-cubin files when available - `ensure_symlink()` maps C++ include paths to the artifact directory Callers (fused_moe, moe_utils) use these two primitives directly. A backward-compatible `get_cubin = get_artifact` alias is kept. Fixes: vllm-project/vllm#38110 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
74c8305 to
4459b55
Compare
|
question: why did this PR touch the cutlass submodule? |
aleozlx
left a comment
aleozlx
left a comment
There was a problem hiding this comment.
lgtm except for the one question above
|
[FAILED] Pipeline #47272725: 7/20 passed |
jimmyzho
left a comment
jimmyzho
left a comment
There was a problem hiding this comment.
lgtm, we could address different hash in a followup
which different hash did you mean? about cutlass or bmm artifacts? |
Different bmm artifacts |
|
oh i think you mean the comment you left earlier.. i have addressed locally. let me push it |
|
i'm making a pass at other chatbot comments to double check... |
|
standby as i test the full workflow locally to double check everything |
|
/bot run |
|
cmds 1 pip install --upgrade setuptools |
|
local tests are looking good, no download observed by the aforementioned approach |
|
disclaimer on the last commit: it says debug but the debug var introduced is harmless and can help future debugging this area. hence leaving it in |
|
i'll send out a .post1 release that cherry picks it to resolve the filed issue, as planned, once this gets merged |
|
irrelevant err on gb200/300 (likely infra error) |
|
[FAILED] Pipeline #47589987: 10/20 passed |
|
AOT builds has timed out |
4 participants
Summary
download_trtllm_headers()stored BMM export headers at a custom path (CUBIN_DIR/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/), butdownload_artifacts()(used when building the flashinfer-cubin wheel) stores them at the canonical artifact hash path. The separateget_file()function only checked the destination path, so it never found the pre-packaged files and always re-downloaded ~17 header files from the network on every startup (×2 for each TP worker).download_trtllm_headers(),get_file(), and renameget_cubin()→get_artifact(). Two unified primitives remain:get_artifact(name, sha256)— fetches any artifact (cubins, headers, checksums, metainfo) into the canonical path, reusing flashinfer-cubin files when availableensure_symlink(link, target)— maps C++ include paths to the artifact directoryget_cubin = get_artifactalias is kept for backward compatibility.Related issue
flashinfer-cubindoes not include all cubins/headers vllm-project/vllm#38110Test plan
get_artifact()local cache (cubins, checksums.txt, metainfo headers, sha256 mismatch) and the BMM header symlink pattern (no download, stale dir replacement, idempotent)🤖 Generated with Claude Code