build: add sccache-backed jit-cache builds and AOT diagnostics - #3205
Merged
Merged
Conversation
…builds Add compiler launcher support (FLASHINFER_NVCC_LAUNCHER, FLASHINFER_CXX_LAUNCHER) to ninja build rules, enabling sccache/ccache wrapping of nvcc and cxx invocations. When unset, behavior is unchanged. Coordinate FLASHINFER_NVCC_THREADS with MAX_JOBS so nvcc parallelizes across gencode targets internally (--threads=N) instead of compiling them sequentially. This better utilizes available cores, especially for builds targeting many SM architectures (e.g., CUDA 12.9 with 7 gencode targets). Pass sccache S3 credentials from GitHub secrets into the Docker build containers for nightly and release workflows. sccache is conditionally installed at build time only when SCCACHE_BUCKET is set; without it, builds behave identically to before. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Apply the same coordinated NVCC_THREADS/MAX_JOBS parallelism and optional sccache integration to task_test_jit_cache_package_build_import.sh, which also builds jit-cache wheels via python -m build --wheel. Moved the parallelism calculation after CUDA arch detection so the thread count is based on the actual number of gencode targets. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Pass S3 sccache credentials through ci/bash.sh into the Docker container for both the spot and on-demand rerun variants of the aot-build-import job. Secrets are not exposed to fork PRs (GitHub Actions default behavior), so this is safe for an open source project. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Instead of requiring AWS credentials (which aren't available to fork PRs), the AOT build import test now always sets up sccache with the public cache bucket using anonymous read-only access (SCCACHE_S3_NO_CREDENTIALS=true). When AWS credentials ARE available (nightly/release builds), sccache operates in read-write mode to populate the cache. This means: - Fork PRs: read-only cache hits, no credentials needed - Nightly/release: read-write, populates the cache for everyone - No secrets passed to pr-test.yml at all Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
EC2 runners are in us-west-2. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
SCCACHE_BUCKET and SCCACHE_REGION are not sensitive — use repository variables (vars.*) instead of secrets so they propagate to fork PR workflows. Only the AWS credentials remain as secrets. Also rename default bucket to flashinfer-build-cache. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
dierksen
requested review from
aleozlx,
bkryu,
cyx-6,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
April 29, 2026 21:10
Contributor
|
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:
📝 WalkthroughWalkthroughAdds jobserver-aware nvcc parallelism and sccache-backed remote compilation caching: exports/control of Changes
Sequence Diagram(s)sequenceDiagram
actor GHA as GitHub Actions
participant Docker as Docker Container
participant Script as Build Script
participant Sccache as sccache Server
participant S3 as S3 Storage
participant Compiler as nvcc/C++ Compiler
GHA->>Docker: Run container, pass env (FLASHINFER_NVCC_THREADS, SCCACHE_BUCKET, SCCACHE_REGION, AWS_*, ...)
Docker->>Script: Execute build_flashinfer_jit_cache_whl.sh
alt SCCACHE_BUCKET set
Script->>Sccache: Install & start sccache
Script->>Sccache: Configure S3 backend (bucket, region, prefix)
Script->>Script: Export FLASHINFER_NVCC_LAUNCHER / FLASHINFER_CXX_LAUNCHER = sccache
end
Script->>Compiler: Invoke compiler (via launcher if set / inherit jobserver if available)
Compiler->>Sccache: Request compile (when launcher used)
Sccache->>S3: Fetch/store cache objects
S3-->>Sccache: Cache hit/miss
Sccache-->>Compiler: Return cached artifact or compiled output
Script->>Sccache: Print sccache stats
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request integrates compiler launchers like sccache into the JIT build system and optimizes build parallelism by coordinating Ninja jobs with internal nvcc threads. Feedback highlights the need to explicitly pass the nvcc_threads variable to Ninja to enable parallel compilation across multiple architectures. Additionally, it is recommended to increase the memory budget per job to mitigate OOM risks and to unify the sccache configuration logic across scripts for better consistency.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
scripts/build_flashinfer_jit_cache_whl.sh (1)
16-32: ClampNVCC_THREADStonprocbefore derivingMAX_JOBS.If a runner has fewer cores than gencode targets, this still oversubscribes a single
nvccprocess. For example, 4 cores and 7 arch targets end up atMAX_JOBS=1andNVCC_THREADS=7, so the “Cap total threads at available CPUs” comment is not actually enforced.Suggested fix
NUM_ARCHS=$(echo "${FLASHINFER_CUDA_ARCH_LIST}" | wc -w) NVCC_THREADS=${FLASHINFER_NVCC_THREADS:-${NUM_ARCHS}} +if (( NVCC_THREADS > NPROC )); then NVCC_THREADS=${NPROC}; fi if (( NVCC_THREADS > 8 )); then NVCC_THREADS=8; fi if (( NVCC_THREADS < 1 )); then NVCC_THREADS=1; fi🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/build_flashinfer_jit_cache_whl.sh` around lines 16 - 32, The NVCC_THREADS value must be clamped to available CPU count before computing MAX_JOBS so a single nvcc process can't be oversubscribed; after you compute and bound NVCC_THREADS (the existing logic around FLASHINFER_NVCC_THREADS and the 1..8 clamp), add a clamp against NPROC (if NVCC_THREADS > NPROC then set NVCC_THREADS=NPROC, ensuring it's still >=1) and only then compute MEM_PER_JOB, MAX_JOBS and TOTAL_THREADS (references: NVCC_THREADS, NPROC, MEM_PER_JOB, MAX_JOBS, TOTAL_THREADS, MEM_AVAILABLE_GB).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@scripts/build_flashinfer_jit_cache_whl.sh`:
- Around line 70-76: The current sccache installation pipes the remote tarball
directly into tar and installs to /usr/local/bin without verifying integrity;
change the flow in scripts/build_flashinfer_jit_cache_whl.sh (and replicate the
same change in scripts/task_test_jit_cache_package_build_import.sh) to download
the tarball to a temporary file (use SCCACHE_VERSION and SCCACHE_ARCH to
construct the URL), also download or embed a pinned checksum for that exact
release, verify the tarball with sha256sum (or another trusted verifier) before
extracting, only extract after the checksum matches, and then move the sccache
binary to /usr/local/bin and clean up; ensure you fail the script with a clear
error if the checksum verification fails.
- Around line 68-89: When SCCACHE_BUCKET is set but AWS credentials are not
available in forked PRs, export SCCACHE_S3_NO_CREDENTIALS=true before
configuring and starting sccache so the runner uses the public read-only S3
cache; update the SCCACHE block (the section that exports SCCACHE_S3_KEY_PREFIX,
SCCACHE_IDLE_TIMEOUT, FLASHINFER_NVCC_LAUNCHER, FLASHINFER_CXX_LAUNCHER and
calls sccache --start-server) to detect missing credentials (e.g.,
AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / AWS_SESSION_TOKEN unset) and in that
case export SCCACHE_S3_NO_CREDENTIALS=true and echo that state so sccache will
operate credentialless for public cache access.
---
Nitpick comments:
In `@scripts/build_flashinfer_jit_cache_whl.sh`:
- Around line 16-32: The NVCC_THREADS value must be clamped to available CPU
count before computing MAX_JOBS so a single nvcc process can't be
oversubscribed; after you compute and bound NVCC_THREADS (the existing logic
around FLASHINFER_NVCC_THREADS and the 1..8 clamp), add a clamp against NPROC
(if NVCC_THREADS > NPROC then set NVCC_THREADS=NPROC, ensuring it's still >=1)
and only then compute MEM_PER_JOB, MAX_JOBS and TOTAL_THREADS (references:
NVCC_THREADS, NPROC, MEM_PER_JOB, MAX_JOBS, TOTAL_THREADS, MEM_AVAILABLE_GB).
🪄 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: 7290836a-da32-42fe-aecf-231956edcaaf
📒 Files selected for processing (5)
.github/workflows/nightly-release.yml.github/workflows/release.ymlflashinfer/jit/cpp_ext.pyscripts/build_flashinfer_jit_cache_whl.shscripts/task_test_jit_cache_package_build_import.sh
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
scripts/build_flashinfer_jit_cache_whl.sh (1)
69-90:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMissing credentialless S3 cache access handling.
This block enables
sccacheunconditionally whenSCCACHE_BUCKETis set, but fork PRs do not receive AWS secrets. Without settingSCCACHE_S3_NO_CREDENTIALS=true, those jobs will fail to use the public read-only cache path.Compare with
scripts/task_test_jit_cache_package_build_import.shwhich correctly handles this case at lines 123-129.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/build_flashinfer_jit_cache_whl.sh` around lines 69 - 90, The sccache setup enables caching whenever SCCACHE_BUCKET is set but doesn't handle credentialless public-read-only access for fork PRs; update the block that configures sccache (the logic setting SCCACHE_S3_KEY_PREFIX, FLASHINFER_NVCC_LAUNCHER, FLASHINFER_CXX_LAUNCHER and starting sccache) to detect when AWS credentials are absent and set SCCACHE_S3_NO_CREDENTIALS=true accordingly (mirror the check used in scripts/task_test_jit_cache_package_build_import.sh), so that when SCCACHE_BUCKET is present but no AWS secrets are available the script uses credentialless S3 access instead of failing.
🤖 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/cpp_ext.py`:
- Around line 183-200: The nvcc jobserver check in should_use_nvcc_jobserver
currently allows CUDA 13.0 but the --jobserver flag requires CUDA 13.1+, so
update the version comparison in should_use_nvcc_jobserver from Version("13.0")
to Version("13.1"); ensure get_nvcc_parallelism_flags continues to call
should_use_nvcc_jobserver(cuda_version) so the --jobserver flag is only appended
for CUDA >= 13.1 (functions: should_use_nvcc_jobserver and
get_nvcc_parallelism_flags).
---
Duplicate comments:
In `@scripts/build_flashinfer_jit_cache_whl.sh`:
- Around line 69-90: The sccache setup enables caching whenever SCCACHE_BUCKET
is set but doesn't handle credentialless public-read-only access for fork PRs;
update the block that configures sccache (the logic setting
SCCACHE_S3_KEY_PREFIX, FLASHINFER_NVCC_LAUNCHER, FLASHINFER_CXX_LAUNCHER and
starting sccache) to detect when AWS credentials are absent and set
SCCACHE_S3_NO_CREDENTIALS=true accordingly (mirror the check used in
scripts/task_test_jit_cache_package_build_import.sh), so that when
SCCACHE_BUCKET is present but no AWS secrets are available the script uses
credentialless S3 access instead of failing.
🪄 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: ac54837e-e8f4-4a0a-b22f-f2cef3c46579
📒 Files selected for processing (7)
.github/workflows/nightly-release.yml.github/workflows/release.ymlflashinfer/jit/core.pyflashinfer/jit/cpp_ext.pyscripts/build_flashinfer_jit_cache_whl.shscripts/task_test_jit_cache_package_build_import.shtests/test_jit_cpp_ext.py
🚧 Files skipped from review as they are similar to previous changes (2)
- .github/workflows/release.yml
- .github/workflows/nightly-release.yml
Contributor
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 (1)
scripts/build_flashinfer_jit_cache_whl.sh (1)
99-101:⚠️ Potential issue | 🟠 MajorAdd
--no-isolationto the wheel build command.This production wheel builder must use the current environment's CUDA/NVSHMEM stack. An isolated build will install the latest nvidia-nvshmem-cu12 from PyPI, creating a version mismatch between the compiled module's device library and the runtime's host library (provided by torch). This causes runtime failures in NVSHMEM-dependent operations.
Suggested change
# Build the wheel using the build module for better isolation echo "Building wheel..." -python -m build --wheel +python -m build --wheel --no-isolation🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/build_flashinfer_jit_cache_whl.sh` around lines 99 - 101, The wheel build invocation "python -m build --wheel" uses an isolated build environment which pulls PyPI dependencies (e.g., nvidia-nvshmem-cu12); change the invocation to disable isolation by adding the --no-isolation flag so the builder uses the current environment's CUDA/NVSHMEM stack (update the "python -m build --wheel" command to "python -m build --wheel --no-isolation").
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@scripts/build_flashinfer_jit_cache_whl.sh`:
- Around line 17-21: The wheel build should pass --no-isolation to avoid
pip-resolving nvidia-nvshmem-cu12; change the python invocation to use "python
-m build --wheel --no-isolation". Also validate FLASHINFER_NVCC_THREADS before
doing arithmetic: parse/convert FLASHINFER_NVCC_THREADS to an integer (fall back
to NUM_ARCHS if empty or non-numeric) and assign that to NVCC_THREADS (or a temp
VAR) before the existing comparisons that use NUM_ARCHS, NVCC_THREADS and NPROC
so the (( ... )) arithmetic never sees non-numeric values.
In `@scripts/task_test_jit_cache_package_build_import.sh`:
- Around line 497-502: Before running the build in the flashinfer-jit-cache
step, remove any existing dist/ artifacts so WHEEL_FILE won't pick up an old
wheel; specifically, add a cleanup that deletes or empties the dist directory
before invoking run_with_aot_memory_monitor "build_flashinfer_jit_cache_wheel"
(python -m build --wheel), then keep the current WHEEL_FILE assignment
(WHEEL_FILE=$(ls dist/*.whl | head -n 1)) so it will select the newly produced
wheel from the fresh dist/ directory.
---
Outside diff comments:
In `@scripts/build_flashinfer_jit_cache_whl.sh`:
- Around line 99-101: The wheel build invocation "python -m build --wheel" uses
an isolated build environment which pulls PyPI dependencies (e.g.,
nvidia-nvshmem-cu12); change the invocation to disable isolation by adding the
--no-isolation flag so the builder uses the current environment's CUDA/NVSHMEM
stack (update the "python -m build --wheel" command to "python -m build --wheel
--no-isolation").
🪄 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: e7783455-3b13-436e-a7f9-80e6cf97d0df
📒 Files selected for processing (2)
scripts/build_flashinfer_jit_cache_whl.shscripts/task_test_jit_cache_package_build_import.sh
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (2)
flashinfer/jit/cpp_ext.py (1)
183-200:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUpdate CUDA version check from 13.0 to 13.1 for
--jobserverflag.The
--jobserverflag was introduced in CUDA 13.1, not 13.0. The current check at line 189 usesVersion("13.0")but should useVersion("13.1")to prevent attempting to use an unsupported flag.🐛 Proposed fix
def should_use_nvcc_jobserver(cuda_version: Optional[Version] = None) -> bool: """Return whether nvcc can safely acquire slots from the inherited jobserver.""" if not sys.platform.startswith("linux"): return False if cuda_version is None: cuda_version = get_cuda_version() - return cuda_version >= Version("13.0") and ( + return cuda_version >= Version("13.1") and ( should_use_ninja_jobserver() or _launcher_uses_sccache() )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/jit/cpp_ext.py` around lines 183 - 200, The CUDA version check in should_use_nvcc_jobserver is wrong: change the minimum required Version from "13.0" to "13.1" so the function only enables the nvcc --jobserver flag for CUDA >= 13.1; update the Version("13.0") literal used in should_use_nvcc_jobserver (which get_nvcc_parallelism_flags relies on) to Version("13.1") so flags.append("--jobserver") is only attempted for supported CUDA versions.scripts/task_test_jit_cache_package_build_import.sh (1)
499-504:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winClean
dist/directory before building wheel.As noted in the past review, if the workspace contains a previous build artifact, line 504 (
ls dist/*.whl | head -n 1) may select an older wheel. Add cleanup before the build step.🔧 Suggested fix
cd flashinfer-jit-cache +rm -rf dist build *.egg-info run_with_aot_memory_monitor "build_flashinfer_jit_cache_wheel" \ python -m build --wheel🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/task_test_jit_cache_package_build_import.sh` around lines 499 - 504, Before running build_flashinfer_jit_cache_wheel, clear previous artifacts so WHEEL_FILE (computed from ls dist/*.whl) won't pick an old wheel: in the cd flashinfer-jit-cache block (where build_flashinfer_jit_cache_wheel and python -m build --wheel are invoked) add a cleanup step to remove the dist directory or delete dist/*.whl (e.g., rm -rf dist || true or rm -f dist/*.whl) immediately before running the build so the subsequent WHEEL_FILE assignment always picks the newly produced wheel.
🧹 Nitpick comments (3)
.github/workflows/pr-test.yml (2)
334-339: 💤 Low valueSame quoting issue for
${DOCKER_IMAGE}.Apply the same fix as noted above for the
aot-build-importjob.🔧 Suggested fix
- bash ci/bash.sh ${DOCKER_IMAGE} \ + bash ci/bash.sh "${DOCKER_IMAGE}" \🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/pr-test.yml around lines 334 - 339, The workflow step invokes bash ci/bash.sh ${DOCKER_IMAGE} ... without quoting the DOCKER_IMAGE variable; update that invocation to quote the variable (use "${DOCKER_IMAGE}") so docker image values with special chars or spaces are handled safely—modify the bash call in the task_test_jit_cache_package_build_import.sh step (same pattern as fixed in the aot-build-import job) to replace ${DOCKER_IMAGE} with "${DOCKER_IMAGE}" while leaving the rest of the environment flags intact.
242-247: 💤 Low valueQuote
${DOCKER_IMAGE}to prevent word splitting.Static analysis flags SC2086 on the unquoted
${DOCKER_IMAGE}variable. While Docker image names typically don't contain spaces, quoting prevents unexpected behavior if the tag contains special characters.🔧 Suggested fix
- bash ci/bash.sh ${DOCKER_IMAGE} \ + bash ci/bash.sh "${DOCKER_IMAGE}" \🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/pr-test.yml around lines 242 - 247, The Docker image variable is unquoted causing shell word-splitting (SC2086); update the call to pass the image as a single quoted argument by changing the invocation of ci/bash.sh to use "${DOCKER_IMAGE}" instead of ${DOCKER_IMAGE} so the bash wrapper receives the full image string intact.scripts/task_test_jit_cache_package_build_import.sh (1)
462-467: Add SHA256 checksum verification for sccache download.The script downloads sccache without verifying the checksum. Mozilla publishes SHA256 checksums for release binaries as separate
.sha256files (e.g.,sccache-v0.9.1-x86_64-unknown-linux-musl.tar.gz.sha256). Adding checksum verification is recommended for supply chain security, especially since this binary will execute compiler commands. Consider downloading and verifying the checksum before extracting the binary.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/task_test_jit_cache_package_build_import.sh` around lines 462 - 467, The sccache download currently pipes the tarball straight to tar without integrity checks; modify the block using SCCACHE_VERSION and SCCACHE_ARCH to first download the tar.gz and its corresponding .sha256 file, verify the tarball with sha256sum (or by computing and comparing the checksum) and exit non‑zero on mismatch, and only after successful verification extract, move and chmod the sccache binary (replacing the existing curl | tar step with explicit download + checksum verification + extraction).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@flashinfer/jit/cpp_ext.py`:
- Around line 183-200: The CUDA version check in should_use_nvcc_jobserver is
wrong: change the minimum required Version from "13.0" to "13.1" so the function
only enables the nvcc --jobserver flag for CUDA >= 13.1; update the
Version("13.0") literal used in should_use_nvcc_jobserver (which
get_nvcc_parallelism_flags relies on) to Version("13.1") so
flags.append("--jobserver") is only attempted for supported CUDA versions.
In `@scripts/task_test_jit_cache_package_build_import.sh`:
- Around line 499-504: Before running build_flashinfer_jit_cache_wheel, clear
previous artifacts so WHEEL_FILE (computed from ls dist/*.whl) won't pick an old
wheel: in the cd flashinfer-jit-cache block (where
build_flashinfer_jit_cache_wheel and python -m build --wheel are invoked) add a
cleanup step to remove the dist directory or delete dist/*.whl (e.g., rm -rf
dist || true or rm -f dist/*.whl) immediately before running the build so the
subsequent WHEEL_FILE assignment always picks the newly produced wheel.
---
Nitpick comments:
In @.github/workflows/pr-test.yml:
- Around line 334-339: The workflow step invokes bash ci/bash.sh ${DOCKER_IMAGE}
... without quoting the DOCKER_IMAGE variable; update that invocation to quote
the variable (use "${DOCKER_IMAGE}") so docker image values with special chars
or spaces are handled safely—modify the bash call in the
task_test_jit_cache_package_build_import.sh step (same pattern as fixed in the
aot-build-import job) to replace ${DOCKER_IMAGE} with "${DOCKER_IMAGE}" while
leaving the rest of the environment flags intact.
- Around line 242-247: The Docker image variable is unquoted causing shell
word-splitting (SC2086); update the call to pass the image as a single quoted
argument by changing the invocation of ci/bash.sh to use "${DOCKER_IMAGE}"
instead of ${DOCKER_IMAGE} so the bash wrapper receives the full image string
intact.
In `@scripts/task_test_jit_cache_package_build_import.sh`:
- Around line 462-467: The sccache download currently pipes the tarball straight
to tar without integrity checks; modify the block using SCCACHE_VERSION and
SCCACHE_ARCH to first download the tar.gz and its corresponding .sha256 file,
verify the tarball with sha256sum (or by computing and comparing the checksum)
and exit non‑zero on mismatch, and only after successful verification extract,
move and chmod the sccache binary (replacing the existing curl | tar step with
explicit download + checksum verification + extraction).
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b81b8d67-ab0e-461d-a9d8-0b08fc19f125
📒 Files selected for processing (5)
.github/workflows/pr-test.ymlflashinfer/jit/core.pyflashinfer/jit/cpp_ext.pyscripts/task_test_jit_cache_package_build_import.shtests/test_jit_cpp_ext.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/jit/core.py
Contributor
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 @.github/workflows/pr-test.yml:
- Around line 242-247: The workflow is passing AWS secrets as CLI args which
exposes them; update the implementation of ci/bash.sh (the script parsing -e
flags) to accept a flag form "-e KEY" with no value and in that case read the
value from the process environment (e.g., getenv(KEY)) instead of expecting an
inline value, and then update the workflow invocation lines that currently pass
"${AWS_ACCESS_KEY_ID}" and "${AWS_SECRET_ACCESS_KEY}" to use the new "-e
AWS_ACCESS_KEY_ID" and "-e AWS_SECRET_ACCESS_KEY" form (or alternatively remove
the inline expansions entirely), ensuring the parsing code in ci/bash.sh handles
both "-e KEY value" and "-e KEY" cases and does not expose secrets on the
command line.
In `@scripts/task_test_jit_cache_package_build_import.sh`:
- Around line 508-514: The credential check currently runs with xtrace enabled
and can leak AWS secrets; wrap the conditional that reads and exports
SCCACHE_S3_NO_CREDENTIALS (the if block referencing AWS_ACCESS_KEY_ID,
AWS_SECRET_ACCESS_KEY, SCCACHE_S3_NO_CREDENTIALS) in a temporary xtrace
disable/restore: turn off tracing (set +x) immediately before the if, perform
the exports/unsets and echo mode, then restore tracing (set -x) afterwards so no
secret values are printed to the trace.
🪄 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: 1a73c83c-7233-4123-aac6-debf6980e74d
📒 Files selected for processing (3)
.github/workflows/pr-test.ymlscripts/build_flashinfer_jit_cache_whl.shscripts/task_test_jit_cache_package_build_import.sh
…c-threads # Conflicts: # scripts/task_test_jit_cache_package_build_import.sh
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
dierksen
changed the title
dierksen
changed the title
Collaborator
Author
|
/bot run |
Collaborator
yongwww
reviewed
May 6, 2026
yongwww
reviewed
May 7, 2026
Move the parallelism computation, sccache installer, and sccache env configuration out of build_flashinfer_jit_cache_whl.sh and task_test_jit_cache_package_build_import.sh into a new scripts/jit_cache_build_common.sh sourced by both. Net -67 lines and both callers now share one source of truth for these helpers. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
dierksen
enabled auto-merge (squash)
May 7, 2026 21:35
aleozlx
pushed a commit
that referenced
this pull request
May 12, 2026
## Description This PR adds sccache-backed compiler caching for the FlashInfer jit-cache/AOT build paths and folds in the AOT memory diagnostics work from #3204. The main pieces are: - Add compiler launcher support in the JIT build generation so `nvcc` and host compiler invocations can be wrapped by `sccache` or another launcher. - Install and configure `sccache` in the release/nightly jit-cache wheel builds and the PR AOT build/import jobs. - Use the shared S3 cache in read-write mode when AWS credentials are available, and fall back to anonymous read-only mode for PR jobs without credentials. - Print `sccache --show-stats` at the end of the relevant jobs so cache hit rates and compile behavior are visible in CI logs. - Keep `FLASHINFER_NVCC_THREADS` supported, but default it back to `1` after CI showed higher nvcc internal threading was slower for this workload. - Increase build-level parallelism through `MAX_JOBS`, bounded by CPU count and an AOT memory budget so the larger CUDA builds do not overrun the runner. - Add AOT memory monitoring/report files around install, wheel build, import/config, and module verification steps to diagnose OOMs and runner shutdowns. - Clean stale jit-cache build artifacts before building and pick the newest wheel from `dist/`. - Add tests for the nvcc flag generation, launcher-compatible depfile/debug flags, and build regeneration behavior. ## Context This PR supersedes #3204. The memory monitor and AOT `MAX_JOBS` safety work from #3204 are included here, but the final script shape is different because it is integrated with the sccache setup and the later threading/concurrency experiments. Merging #3204 separately would still create a conflict in `scripts/task_test_jit_cache_package_build_import.sh`. The current direction is to use sccache for the big rebuild/retry win while keeping nvcc internal threading conservative. The PR still allows `FLASHINFER_NVCC_THREADS` to be overridden for experiments, but the default path lets ninja drive parallelism through `MAX_JOBS`. ## CI Notes The latest PR Test run completed successfully after the infrastructure-triggered AOT rerun path. The Release workflow's jit-cache jobs completed and showed useful sccache hit rates; the remaining Release failure was in `build-flashinfer-cubin` while downloading cubins, which appears separate from the sccache/AOT build changes. ## Related Supersedes #3204. ## Performance comparison This looks at the mean and median time to complete each AOT build type over the last 48 hours compared with the most recent 2 builds on this branch. | job | 48h mean | 48h median | [#4214](https://github.com/flashinfer-ai/flashinfer/actions/runs/25402485487) | #4214 vs median | [#4284](https://github.com/flashinfer-ai/flashinfer/actions/runs/25457504611) | #4284 vs median | |---|---:|---:|---:|---:|---:|---:| | x64 cu126 | 1:09:34 | 51:55 | 42:28 | 18.2% faster | 31:36 | 39.1% faster | | x64 cu128 | 1:45:24 | 1:24:04 | 1:09:46 | 17.0% faster | 45:04 | 46.4% faster | | x64 cu129 | 2:36:17 | 2:25:29 | 2:12:39 | 8.8% faster | 59:50 | 58.9% faster | | x64 cu130 | 2:27:16 | 2:04:49 | 1:48:57 | 12.7% faster | 59:34 | 52.3% faster | | arm64 cu126 | 1:23:52 | 1:30:43 | 1:13:56 | 18.5% faster | 50:43 | 44.1% faster | | arm64 cu128 | 2:12:08 | 2:25:56 | 2:04:47 | 14.5% faster | 1:12:37 | 50.2% faster | | arm64 cu129 | 3:45:38 | 4:05:35 | 3:38:02 | 11.2% faster | 1:31:15 | 62.8% faster | | arm64 cu130 | 3:34:21 | 3:53:18 | 3:21:49 | 13.5% faster | 1:34:48 | 59.4% faster | --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
5 tasks
3 tasks
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
3 participants
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.
Description
This PR adds sccache-backed compiler caching for the FlashInfer jit-cache/AOT build paths and folds in the AOT memory diagnostics work from #3204.
The main pieces are:
nvccand host compiler invocations can be wrapped bysccacheor another launcher.sccachein the release/nightly jit-cache wheel builds and the PR AOT build/import jobs.sccache --show-statsat the end of the relevant jobs so cache hit rates and compile behavior are visible in CI logs.FLASHINFER_NVCC_THREADSsupported, but default it back to1after CI showed higher nvcc internal threading was slower for this workload.MAX_JOBS, bounded by CPU count and an AOT memory budget so the larger CUDA builds do not overrun the runner.dist/.Context
This PR supersedes #3204. The memory monitor and AOT
MAX_JOBSsafety work from #3204 are included here, but the final script shape is different because it is integrated with the sccache setup and the later threading/concurrency experiments. Merging #3204 separately would still create a conflict inscripts/task_test_jit_cache_package_build_import.sh.The current direction is to use sccache for the big rebuild/retry win while keeping nvcc internal threading conservative. The PR still allows
FLASHINFER_NVCC_THREADSto be overridden for experiments, but the default path lets ninja drive parallelism throughMAX_JOBS.CI Notes
The latest PR Test run completed successfully after the infrastructure-triggered AOT rerun path. The Release workflow's jit-cache jobs completed and showed useful sccache hit rates; the remaining Release failure was in
build-flashinfer-cubinwhile downloading cubins, which appears separate from the sccache/AOT build changes.Related
Supersedes #3204.
Performance comparison
This looks at the mean and median time to complete each AOT build type over the last 48 hours compared with the most recent 2 builds on this branch.