Conversation
The cu130 flashinfer-jit-cache wheel grew to 2.0 GB and started failing to upload as a GitHub Release asset (per-asset 2 GiB limit; see flashinfer-ai#3257). Each new SM target appended ~150-200 MB compressed to every wheel, and cu130 carries 8 (sm75/80/89/90a/100a/103a/110a/120f). Split each (CUDA, CPU-arch) wheel into three by GPU SM family: - sm9x - Ampere/Ada/Hopper (<= sm90a) - sm10x - Datacenter Blackwell (sm100a/103a/110a) - sm12x - Consumer Blackwell (sm120f, future sm121a) Same package name everywhere; the family is encoded in the PEP 440 local-version, so wheels resolve as e.g. 'flashinfer-jit-cache== 0.6.11+cu130.sm10x'. Existing 'pip install flashinfer-jit-cache' still works once the right pin is given. Driven by a new 'flashinfer install-jit-cache-wheel' subcommand that detects FlashInfer version, CUDA version, and GPU compute capability (via torch.cuda.get_device_capability) and runs the matching pip install. Honors --cuda-version, --sm-family, --nightly, --dry-run. Modeled on the CLI scaffolding from flashinfer-ai#3142 with the family dimension added. Build side: 'FLASHINFER_JIT_CACHE_SM_FAMILY' env var, when set, filters 'FLASHINFER_CUDA_ARCH_LIST' to the family's archs and appends '.<family>' to the local-version suffix. Release / nightly workflows gain an 'sm_family' matrix dimension; the upload-to-release loop iterates over all three families. The wheel-index regex accepts the new local-version shape and remains compatible with the legacy '+cuXY' format. Closes flashinfer-ai#3257 Co-Authored-By: Claude Opus 4.7 (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
May 7, 2026 23:25
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 SM-family-aware JIT-cache wheel distribution: utilities to map/filter CUDA arches by SM family, build backend and metadata changes, a new install CLI with autodetection, CI matrix/artifact updates, wheel-index parsing, docs updates, and corresponding tests. ChangesSM Family JIT Cache Wheels
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 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)
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 introduces a new CLI command, flashinfer install-jit-cache-wheel, to automate the installation of flashinfer-jit-cache wheels by autodetecting the CUDA version and GPU SM family. This update supports a new distribution model where wheels are split by SM family to comply with GitHub's asset size limits. Feedback from the review suggests improving the robustness of CUDA architecture parsing to handle various string formats and refactoring duplicated SM family logic into a common utility to enhance maintainability.
Mirrors the per-family split in release/nightly so PR CI actually exercises the per-family build path. Was previously running one job per (cuda, arch) which still built every arch; now runs three jobs per (cuda, arch) — one per SM family — each compiling only its family's archs. - pr-test.yml: 'aot-build-import' and 'aot-build-import-rerun' gain 'sm_family: [sm9x, sm10x, sm12x]'. cu126 is excluded for sm10x and sm12x because that toolkit only supports archs <= sm90. The rerun matrix builder mirrors the same exclude. FLASHINFER_JIT_CACHE_SM_FAMILY is forwarded into the test container via ci/bash.sh's '-e' flag. - task_test_jit_cache_package_build_import.sh: when FLASHINFER_JIT_CACHE_SM_FAMILY is set, filter FLASHINFER_CUDA_ARCH_LIST to that family's archs before running the wheel build and verify_all_modules_compiled.py. The build-side filter in build_backend.py mutates os.environ inside its own process only, so doing it once in the parent shell ensures both subprocesses see the same arch list. Also fix black formatting flagged by pre-commit on PR flashinfer-ai#3265: - build_backend.py: rewrite SM_FAMILIES lambdas as named functions to avoid black's awkward multi-line break of '<' chained comparisons. - __main__.py: collapse a ClickException to single-line per black's preference. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
![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: 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-cache/build_backend.py`:
- Around line 38-94: Run the project formatter (e.g. ruff format or pre-commit
run --all-files) and commit the resulting changes so the SM_FAMILIES dict and
the multi-line print in _apply_sm_family_filter are formatted to satisfy ruff;
specifically reformat the SM_FAMILIES declaration and the print(...) call in
_apply_sm_family_filter (and any other affected lines) and push the reformatted
file so CI passes.
In `@flashinfer/__main__.py`:
- Around line 267-279: Re-run the project's formatter (ruff format) to apply the
canonical formatting for the Click exception lines in the CUDA-version parsing
block: ensure the click.ClickException(...) call around the InvalidVersion
exception handling and the earlier validation (the calls that raise
click.ClickException when normalized startswith "cu" and in the except block
that wraps InvalidVersion) are formatted according to ruff so the pre-commit
check passes; after formatting, stage and commit the changes.
- Around line 350-403: The current install_jit_cache_wheel_cmd builds an exact
pinned requirement from resolved_flashinfer_version which breaks when --nightly
points at nightly index but the installed __version__ is a stable release;
modify install_jit_cache_wheel_cmd to detect nightly and, if nightly is True and
resolved_flashinfer_version is a release (no "dev" or "+"), construct a range
requirement instead of an exact pin (e.g.
"flashinfer-jit-cache>={base},<{next_major_or_minor}") by parsing
resolved_flashinfer_version with packaging.version to compute the next version
bound, or alternatively call a new flag-aware helper (update
_build_jit_cache_requirement or add _build_jit_cache_requirement_for_nightly)
that returns the looser requirement when nightly is set; ensure the printed
requirement and pip args use this new requirement variable.
🪄 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: 7e3e084a-a262-4e22-b1ba-8cc3a970463e
📒 Files selected for processing (7)
.github/workflows/nightly-release.yml.github/workflows/release.ymlREADME.mddocs/installation.rstflashinfer-jit-cache/build_backend.pyflashinfer/__main__.pyscripts/update_whl_index.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 @.github/workflows/pr-test.yml:
- Around line 244-247: The run step invokes ci/bash.sh with an unquoted
${DOCKER_IMAGE}, which triggers SC2086 (word-splitting); update the command to
quote the variable as "$DOCKER_IMAGE" in the invocation (e.g., change ci/bash.sh
${DOCKER_IMAGE} --no-gpu ... to ci/bash.sh "$DOCKER_IMAGE" --no-gpu ...), and
make the same change in the equivalent rerun "Run Test" step that calls the same
command line so both occurrences use "$DOCKER_IMAGE".
🪄 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: 4f367040-09e4-47ff-8782-3ddd358e6b36
📒 Files selected for processing (4)
.github/workflows/pr-test.ymlflashinfer-jit-cache/build_backend.pyflashinfer/__main__.pyscripts/task_test_jit_cache_package_build_import.sh
The Release workflow has a 'pull_request: paths: .github/workflows/
release.yml' trigger that runs the build jobs in dry-run mode whenever
release.yml changes. Its checkout used:
ref: ${{ github.event_name == 'pull_request' && github.head_ref || inputs.tag }}
For a fork PR, github.head_ref resolves to a branch that doesn't exist
on flashinfer-ai/flashinfer (because actions/checkout defaults
'repository:' to the workflow's repo). 'git fetch' fails three times,
and the setup job dies before any actual build work runs.
The bug has been latent since flashinfer-ai#1910 (2025-10-10), where the trigger
and the buggy checkout were introduced together. It only fires on
fork-PRs that touch release.yml; PRs from branches on the main repo
work fine because the default 'repository:' already matches.
Fix all four affected checkouts to set 'repository:' explicitly to
the PR head's repo and pin to head.sha (which is also stable across
re-pushes during the run). On workflow_dispatch the existing
'inputs.tag' path is unchanged.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
/bot run |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/env.py`:
- Around line 126-127: The compatibility check currently uses
_public_package_version(flashinfer_version) !=
_public_package_version(flashinfer_jit_cache_version) but strips the
local-version suffix so different SM-family suffixes (e.g. .sm9x vs .sm12x) are
ignored; update the logic to, when CUDA is available, extract the sm* suffix
from flashinfer_jit_cache_version (e.g. via a small regex on the local-version
segment) and compare it to the detected device family (use your CUDA detection
helper / device-family variable); if the sm suffix is present and does not match
the detected device family, raise the same incompatibility error (or fail fast)
instead of proceeding, while falling back to the existing
_public_package_version check for non-CUDA cases.
🪄 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: 20101ea1-0e1a-416e-84cc-231857008350
📒 Files selected for processing (7)
.github/workflows/nightly-release.yml.github/workflows/pr-test.yml.github/workflows/release.ymlflashinfer/__main__.pyflashinfer/jit/env.pyscripts/task_test_jit_cache_package_build_import.shtests/cli/test_cli_cmds.py
🚧 Files skipped from review as they are similar to previous changes (4)
- scripts/task_test_jit_cache_package_build_import.sh
- .github/workflows/nightly-release.yml
- .github/workflows/pr-test.yml
- .github/workflows/release.yml
Collaborator
kahyunnam
reviewed
May 8, 2026
Member
There was a problem hiding this comment.
re: the 3 GPU SM families in the PR description, I think @aleozlx mentioned earlier in the thread that each device typically requires 8.0 plus their native arch -- should we add sm80a compialtion to sm10x and sm12x subwheels as well?
added 4 commits
May 8, 2026 15:22
# Conflicts: # scripts/task_test_jit_cache_package_build_import.sh
Member
|
|
yongwww
reviewed
May 14, 2026
Collaborator
Author
Yeah, that one in particular does rub me the wrong way, but if I understand correctly, SM110 (previously SM101) maps closer to the datacenter Blackwell chips than to the SM120f chips. That said, maybe we should split things up more to align better with different systems? Using the table from https://developer.nvidia.com/cuda/gpus and our support matrix from https://github.com/flashinfer-ai/flashinfer/#gpu-support, we could probably split things up like: SM75-89: x86_64 and aarch64 Thor stands out as relatively unique, and this would help keep the binary size lower for all of them. |
Collaborator
Author
|
Looking at the breakdown in jit-cache wheel sizes: For cu130:
Currently in this PR:
And the original:
Another thing that changes is that if we align with PyTorch and shift to 12.6/13.0/13.2, we might be able to drop SM100+ on 12.x and only provide those for 13.x. |
Collaborator
Author
|
Questions from meeting today:
|
mgoin
reviewed
May 21, 2026
Comment on lines
+105
to
+107
| # JIT cache: autodetect CUDA + GPU SM family and run the matching pip install. | ||
| # Use --dry-run to preview, --sm-family / --cuda-version to override. | ||
| flashinfer install-jit-cache-wheel |
Contributor
There was a problem hiding this comment.
Can you make sure to test this within regular python venvs and also with uv venvs? I commonly use uv pip install so I forget if there sometimes can be issues mixing pip install and uv pip install within the same environment
Comment on lines
+25
to
+29
| SM_FAMILY_BASE_ARCHS = { | ||
| "sm10x": ("8.0",), | ||
| "sm110": ("8.0",), | ||
| "sm12x": ("8.0",), | ||
| } |
Contributor
There was a problem hiding this comment.
Why are these all 8.0 for base arches?
Collaborator
Author
There was a problem hiding this comment.
8.0 is needed in all of the wheels as a baseline (IIUC, @aleozlx can clarify), and this is just how it's currently encoding that for the 10.0+ wheels. We could also change to just have everything explicitly list 8.0 in each of them.
Contributor
|
Making a note from offline discussion that we will still need something like |
Collaborator
Author
This might require some additional planning to implement. With the current setup, I'm not sure if the .so's would resolve properly if we download them separately and then try to access them. What we might need to do instead is parcel out the existing wheel structure into shards that are under the 2GB release artifact limit enforced by GitHub, then reconstruct that after downloading the shards, only when using the Any thoughts on this implementation or other suggestions? Right now we can't release at all due to the GitHub artifact cap, so we're kind of stuck until we change something. |
5 tasks
aleozlx
pushed a commit
that referenced
this pull request
Sep 10, 2026
## 📌 Description This PR turns `flashinfer-jit-cache` into a small top-level shim backed by independently installable, architecture-specific provider wheels. It retains legacy monolithic cache discovery as a fallback while adding: - provider package metadata and entry-point discovery - exact CUDA architecture matching, including suffixed targets such as `sm90a`, `sm120f`, and `sm121a` - a default installation mode whose shim dependencies name the complete published provider set - a minimal installation mode that selects only an explicitly requested or locally detected provider, with no implicit SM80 baseline - provider wheelhouse build, binary inventory, installation, and JIT-disabled GPU smoke tooling - AOT capability gating so unsupported modules, currently BGMV MoE on SM80, are omitted from that provider rather than failing its build The provider matrix is intentionally independent of the size-pruned monolithic wheel matrix. Architecture-specific wheels may therefore preserve useful native targets such as SM121a without adding them back to every monolithic wheel. Native providers are SASS-only; the shim does not rely on PTX or an SM80 compatibility baseline. ## 🔍 Related Issues - Follow-up design to #3265 - Incorporates current-main behavior from #4469, #4527, #4682, #4711, #4757, and #4760 ## 🧪 Validation ### Provider canaries Two end-to-end CUDA 13.0 canaries have exercised provider build, strict binary inspection, shim/provider installation, provider discovery, top-level import without CUTLASS DSL, and a JIT-disabled `silu_and_mul` numerical smoke: | System | Target | Provider result | | --- | --- | --- | | x86_64 A100 test system | A100 / SM80 | 113.8 MiB; every CUDA-bearing module SM80-only; zero PTX | | DGX Spark | aarch64 GB10 / SM121a | 173.0 MiB; 558 modules; every CUDA-bearing module SM121a-only; zero PTX | The x86_64 A100 artifact was built before the final BGMV capability correction and contained 207 modules. Final-branch source-level SM80 AOT enumeration produces 206 modules with `bgmv_moe` absent, and the associated focused suite passed 30 tests. A final-head SM80 provider wheel still needs to be rebuilt to confirm that exact packaged inventory. The Spark provider could load its packaged BGMV module, but a BGMV numerical invocation exceeded the device's dynamic shared-memory limit (approximately 216 KiB requested versus approximately 101 KiB available). That is a pre-existing kernel/runtime limitation and is outside this packaging change; the independent `silu_and_mul` provider smoke passed. ### Current branch and CI - CodeRabbit passes and all inline review threads are resolved. - Pre-commit, documentation, and public API/documentation checks pass. - The existing monolithic release workflow passes for cu129, cu130, and cu134 on both x86_64 and aarch64. These jobs validate backward compatibility, not provider-wheel publication. - The manually authorized [full PR test run](https://github.com/flashinfer-ai/flashinfer/actions/runs/33890607699) passes all four cu129/cu130 x64/arm64 AOT build-import jobs, all five A10G JIT shards, and the T4 JIT job. The H100 JIT job is still running as of September 4, 2026. - After the latest review fixes, the focused provider suite passed on an x86_64 A100 test system in a disposable container: 25 passed, 1 warning. `pre-commit`, `bash -n`, and `git diff --check` also pass. The branch is currently mergeable. It will be rebased onto `main` once more before merge; intervening main-branch changes reviewed so far do not alter the provider packaging or AOT capability implementation. ## Remaining Validation Before enabling provider publication in a release workflow: - let the current H100 PR test finish and address any real failure - rebuild the final-head SM80 provider and repeat strict inventory, install, and GPU smoke validation - build and inspect a real cu134 provider artifact; current cu134 CI covers only the legacy wheel and static configuration paths - validate default all-provider installation plus minimal auto-detected and explicit-target installation against a multi-provider wheelhouse - teach `scripts/update_whl_index.py` to recognize provider distribution names - add an artifact-only shadow provider matrix to nightly/release automation and collect size, build-time, homogeneous GPU, and heterogeneous GPU results before changing the public release format ## Reviewer Notes The main policy question is the explicit provider coverage matrix for each CUDA and CPU architecture. The current implementation makes the conservative choices: exact target matching, literal shim dependencies, no closest-lower-architecture inference, no implicit SM80 provider, and normal JIT compilation when no compatible AOT provider is available. This is ready for human review of the package contract, installation UX, and release shape. The remaining items above are release-enablement validation rather than evidence that the architectural split itself has not been exercised. ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added experimental architecture-specific JIT-cache provider wheels with automatic discovery and legacy fallback. * Added `install-jit-cache-wheel` options for full or minimal installation and CUDA architecture selection. * Added tools to build, validate, and smoke-test provider packages. * **Bug Fixes** * Improved architecture-aware AOT module selection and generation. * Prevented incompatible architectures from being selected automatically. * **Documentation** * Documented provider-wheel configuration, installation modes, and experimental build options. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
11 tasks
dierksen
added a commit
that referenced
this pull request
Sep 15, 2026
## 📌 Description This PR enables the architecture-specific JIT-cache provider wheels introduced by #4514 and removes the monolithic JIT-cache wheel mode from release, build, package, and runtime paths. - Adds explicit provider inventories per CUDA and CPU platform in `ci/cuda-versions.json`. - Derives the individual provider build matrix and the CUDA/CPU shim matrix from that single inventory. - Builds each binary provider, builds one platform-specific `flashinfer-jit-cache` shim per CUDA/CPU pair, and validates the assembled set before publication. - Builds the patched CUDA 13.4 `sccache` binary once per CPU architecture and shares it across the corresponding provider builds. - Validates provider SASS using CUDA binary-compatibility rules while continuing to reject PTX and targets that cannot execute on the declared architecture. - Propagates the selected architecture flags through NVCC device linking so linked modules do not acquire the toolkit's default `sm52` or `sm75` cubin. - Selects the closest compatible installed provider at runtime, with exact targets preferred over family targets and family targets preferred over unsuffixed baselines. - Requires each provider's full CUDA-specific version to match the installed shim, preventing a provider such as `+cu134` from being loaded by a `+cu130` shim. - Supports heterogeneous processes by resolving the best provider independently for every visible CUDA architecture, loading all selected copies of a module, and dispatching each call according to its CUDA tensor devices. - Keeps partial provider inventories usable: a module present only in the SM103 provider can run on SM103 while an uncovered SM120 call falls back to the normal multi-architecture JIT module. - Registers the cubin callback for every selected provider library instead of only the primary path. - Makes minimal installation discover the provider inventory from the exact published shim wheel instead of maintaining a second architecture table. - Drops separate SM86 provider builds because the unsuffixed SM80 provider is forward compatible within compute-capability major 8. - Teaches wheel-index generation to recognize provider distribution names. - Removes monolithic wheel compilation, artifact publication, package metadata, runtime discovery, and workflow selection paths. - Keeps reusable provider validation independent from the disposable one-off wheelhouse tooling. Pull request runs exercise the complete split provider build, shim assembly, installation, and artifact-validation path without publishing artifacts. Once merged, scheduled nightlies and release dispatches build and publish only the split provider wheel set. Rollback requires reverting to a workflow revision that still builds monolithic wheels; there is no runtime format switch. ## 🔍 Related Issues - Follow-up to #4514 - Design follow-up to #3265 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up pre-commit, see the [pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). Completed on the current revision: - commit-time pre-commit checks for all changed files, including Ruff and mypy - `python3 ci/validate_cuda_versions.py` - `pytest tests/test_env.py -q --noconftest` (`6 passed`), including heterogeneous provider resolution and nested tensor-device discovery - Python byte-compilation and `git diff --check` Completed on an earlier revision of this PR: - `pytest -q --confcutdir=tests/jit tests/jit/test_jit_cache_release.py` (`6 passed` in a lightweight local environment) - provider, release, and CLI suites in a CUDA container on an ARM64 SM121 test system (`86 passed`) - focused NVCC device-link regression test and a direct `cuobjdump --list-elf` check of an affected `sm103a` module (only `sm103a` cubins present) The complete provider build matrix and assembled-wheel installation checks are exercised by this PR's release dry run. Actual same-process dispatch across an SM103 and an SM120 GPU remains a hardware validation gap. The updated canary supports that test as one invocation with `--provider sm103a --provider sm120f` and JIT disabled. ## 🔬 Experimental Track <!-- Only for PRs submitted under the experimental policy (CONTRIBUTING.md → "Experimental APIs and Backends"). Leave this section untouched for normal PRs. --> - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows. ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` ## Reviewer Notes Please focus on the provider inventory policy, CUDA binary-compatibility rules, provider-only workflow fan-in, platform-specific shim tags and dependencies, artifact assembly, removal of monolithic paths, and the heterogeneous AOT dispatcher's tensor-device selection and JIT fallback behavior. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Releases and nightly builds now provide architecture-specific JIT-cache provider wheels alongside a shim wheel. - Provider wheels include validation reports and SHA-256 checksums. - Installed providers are discovered and validated automatically. - CUDA provider architecture coverage has expanded. - **Breaking Changes** - Legacy monolithic JIT-cache wheel builds and format selection are no longer supported. - **Documentation** - Updated JIT-cache documentation to describe provider-based releases and validation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Vincent Tombari <Vinnie6167@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
passtoor-agi
pushed a commit
to passtoor-agi/flashinfer
that referenced
this pull request
Sep 15, 2026
## 📌 Description This PR enables the architecture-specific JIT-cache provider wheels introduced by flashinfer-ai#4514 and removes the monolithic JIT-cache wheel mode from release, build, package, and runtime paths. - Adds explicit provider inventories per CUDA and CPU platform in `ci/cuda-versions.json`. - Derives the individual provider build matrix and the CUDA/CPU shim matrix from that single inventory. - Builds each binary provider, builds one platform-specific `flashinfer-jit-cache` shim per CUDA/CPU pair, and validates the assembled set before publication. - Builds the patched CUDA 13.4 `sccache` binary once per CPU architecture and shares it across the corresponding provider builds. - Validates provider SASS using CUDA binary-compatibility rules while continuing to reject PTX and targets that cannot execute on the declared architecture. - Propagates the selected architecture flags through NVCC device linking so linked modules do not acquire the toolkit's default `sm52` or `sm75` cubin. - Selects the closest compatible installed provider at runtime, with exact targets preferred over family targets and family targets preferred over unsuffixed baselines. - Requires each provider's full CUDA-specific version to match the installed shim, preventing a provider such as `+cu134` from being loaded by a `+cu130` shim. - Supports heterogeneous processes by resolving the best provider independently for every visible CUDA architecture, loading all selected copies of a module, and dispatching each call according to its CUDA tensor devices. - Keeps partial provider inventories usable: a module present only in the SM103 provider can run on SM103 while an uncovered SM120 call falls back to the normal multi-architecture JIT module. - Registers the cubin callback for every selected provider library instead of only the primary path. - Makes minimal installation discover the provider inventory from the exact published shim wheel instead of maintaining a second architecture table. - Drops separate SM86 provider builds because the unsuffixed SM80 provider is forward compatible within compute-capability major 8. - Teaches wheel-index generation to recognize provider distribution names. - Removes monolithic wheel compilation, artifact publication, package metadata, runtime discovery, and workflow selection paths. - Keeps reusable provider validation independent from the disposable one-off wheelhouse tooling. Pull request runs exercise the complete split provider build, shim assembly, installation, and artifact-validation path without publishing artifacts. Once merged, scheduled nightlies and release dispatches build and publish only the split provider wheel set. Rollback requires reverting to a workflow revision that still builds monolithic wheels; there is no runtime format switch. ## 🔍 Related Issues - Follow-up to flashinfer-ai#4514 - Design follow-up to flashinfer-ai#3265 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up pre-commit, see the [pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). Completed on the current revision: - commit-time pre-commit checks for all changed files, including Ruff and mypy - `python3 ci/validate_cuda_versions.py` - `pytest tests/test_env.py -q --noconftest` (`6 passed`), including heterogeneous provider resolution and nested tensor-device discovery - Python byte-compilation and `git diff --check` Completed on an earlier revision of this PR: - `pytest -q --confcutdir=tests/jit tests/jit/test_jit_cache_release.py` (`6 passed` in a lightweight local environment) - provider, release, and CLI suites in a CUDA container on an ARM64 SM121 test system (`86 passed`) - focused NVCC device-link regression test and a direct `cuobjdump --list-elf` check of an affected `sm103a` module (only `sm103a` cubins present) The complete provider build matrix and assembled-wheel installation checks are exercised by this PR's release dry run. Actual same-process dispatch across an SM103 and an SM120 GPU remains a hardware validation gap. The updated canary supports that test as one invocation with `--provider sm103a --provider sm120f` and JIT disabled. ## 🔬 Experimental Track <!-- Only for PRs submitted under the experimental policy (CONTRIBUTING.md → "Experimental APIs and Backends"). Leave this section untouched for normal PRs. --> - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows. ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` ## Reviewer Notes Please focus on the provider inventory policy, CUDA binary-compatibility rules, provider-only workflow fan-in, platform-specific shim tags and dependencies, artifact assembly, removal of monolithic paths, and the heterogeneous AOT dispatcher's tensor-device selection and JIT fallback behavior. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Releases and nightly builds now provide architecture-specific JIT-cache provider wheels alongside a shim wheel. - Provider wheels include validation reports and SHA-256 checksums. - Installed providers are discovered and validated automatically. - CUDA provider architecture coverage has expanded. - **Breaking Changes** - Legacy monolithic JIT-cache wheel builds and format selection are no longer supported. - **Documentation** - Updated JIT-cache documentation to describe provider-based releases and validation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Vincent Tombari <Vinnie6167@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
5 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.
The cu130 flashinfer-jit-cache wheel grew to 2.0 GB and started failing to upload as a GitHub Release asset (per-asset 2 GiB limit; see #3257). Each new SM target appended ~150-200 MB compressed to every wheel, and cu130 carries 8 (sm75/80/89/90a/100a/103a/110a/120f).
Split each (CUDA, CPU-arch) wheel into three by GPU SM family:
Same package name everywhere; the family is encoded in the PEP 440 local-version, so wheels resolve as e.g. 'flashinfer-jit-cache== 0.6.11+cu130.sm10x'. Existing 'pip install flashinfer-jit-cache' still works once the right pin is given.
Driven by a new 'flashinfer install-jit-cache-wheel' subcommand that detects FlashInfer version, CUDA version, and GPU compute capability (via torch.cuda.get_device_capability) and runs the matching pip install. Honors --cuda-version, --sm-family, --nightly, --dry-run. Modeled on the CLI scaffolding from #3142 with the family dimension added.
Build side: 'FLASHINFER_JIT_CACHE_SM_FAMILY' env var, when set, filters 'FLASHINFER_CUDA_ARCH_LIST' to the family's archs and appends '.' to the local-version suffix. Release / nightly workflows gain an 'sm_family' matrix dimension; the upload-to-release loop iterates over all three families. The wheel-index regex accepts the new local-version shape and remains compatible with the legacy '+cuXY' format.
Closes #3257
Closes #3033
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Testing Done
Local DGX Spark validation:
aarch64with NVIDIA GB10, compute capability12.1; PyTorch sees CUDA13.0and reports the device as(12, 1).flashinfer install-jit-cache-wheel --dry-runfrom the source checkout. It now resolves fromversion.txtwhen package metadata is0.0.0+unknown, detects CUDA13.0, selectssm12x, and printsflashinfer-jit-cache==0.6.11+cu130.sm12x.python -m pip, so the CLI now falls back touv pip install --python ....https://flashinfer.ai/whl/cu130; it reached the resolver cleanly and failed only because0.6.11+cu130.sm12xis not yet published.jit-cache-cu129-aarch64-sm12xartifact from Release run25528598737, served it through a local simple index, installed it with explicit--cuda-version cu129 --sm-family sm12x --index-url ..., importedflashinfer_jit_cache, verifiedFLASHINFER_AOT_DIRpoints at the installed package cache, then uninstalled it because the local machine is CUDA 13.0.jit-cache-cu130-aarch64-sm12xartifact is absent because that job failed while downloading build dependencies (IncompleteReadfornvidia_cublas), not because of the CLI install path.Merge/conflict validation:
upstream/mainand resolved conflicts inpr-test.yml,release.yml, andnightly-release.yml, preserving bothFLASHINFER_JIT_CACHE_SM_FAMILYforwarding and upstream sccache/NVCC env forwarding.yaml.safe_load.git diff --cached --check.python -m py_compile flashinfer/__main__.py flashinfer/jit/env.py tests/cli/test_cli_cmds.py.python -m pytest tests/cli/test_cli_cmds.py -q(18 passed).Review-comment follow-up:
90,100, and120in the SM-family filter, and moved the duplicated SM-family helpers intobuild_utils.pyfor reuse by both the CLI and jit-cache build backend.--nightlyconcern by rejecting nightly installs when the resolved FlashInfer version is not a dev release; dev versions still exact-pin the matching SM-family wheel, e.g.flashinfer-jit-cache==0.6.11.dev20260508+cu130.sm12x.python -m py_compile build_utils.py flashinfer/__main__.py flashinfer-jit-cache/build_backend.py tests/cli/test_cli_cmds.py.python -m pytest tests/cli/test_cli_cmds.py -q(21 passed, with the expected PyTorch GB10 capability warning from this host's torch build).git diff --check.flashinfer install-jit-cache-wheel --cuda-version cu130 --sm-family sm12x --dry-run, which resolvesflashinfer-jit-cache==0.6.11+cu130.sm12xand theuv pip install --python ...command.0.6.11now fails early with the new explanatory error, while explicit0.6.11.dev20260508resolvesflashinfer-jit-cache==0.6.11.dev20260508+cu130.sm12xagainsthttps://flashinfer.ai/whl/nightly/cu130with--pre.Human feedback follow-up:
flashinfer install-jit-cache-wheelautodetection to inspect every visible CUDA device instead of only device 0. It selects a wheel only when the visible GPUs are covered by one jit-cache SM-family wheel, and otherwise fails with guidance to pass--sm-familyor build from source with an explicitFLASHINFER_CUDA_ARCH_LIST.sm80base arch plus native Blackwell archs. The build-side family filter now keeps/adds8.0forsm10xandsm12xonly when a native arch for that family is present.sm12xdefault arch lists on12.0f; the family-specificsm120ftarget covers DGX Spark / GB10 (sm121) without adding an exact12.1atarget by default.flashinfer-jit-cachelocal-version SM suffixes. On CUDA hosts, a wrong-family installed wheel now fails fast; on this DGX Spark,0.6.11+cu130.sm12xvalidates and0.6.11+cu130.sm9xfails with an expected-family error.python -m py_compile build_utils.py flashinfer/__main__.py flashinfer/jit/env.py flashinfer-jit-cache/build_backend.py tests/cli/test_cli_cmds.py.python -m pytest tests/cli/test_cli_cmds.py -q(27 passed, with the expected PyTorch GB10 capability warning from this host's torch build).uvx ruff check ...anduvx ruff format --check ...over the touched Python files.git diff --check,bash -n scripts/task_test_jit_cache_package_build_import.sh, and parsed the touched workflow YAML files withyaml.safe_load.sm10x: 8.0 10.0a 10.3a 11.0aandsm12x: 8.0 12.0ffor the CUDA 13.0 release arch list.sm12xand resolvesflashinfer-jit-cache==0.6.11+cu130.sm12xusing theuv pip install --python ...fallback.SM121 target cleanup:
12.1aadditions from release/nightly/default jit-cache arch lists and docs;sm12xnow defaults to8.0 12.0f.12.1asupport in the parser/filter if a user supplies it manually, but release artifacts no longer build it by default.python -m pytest tests/cli/test_cli_cmds.py -q(27 passed),uvx ruff check tests/cli/test_cli_cmds.py,uvx ruff format --check tests/cli/test_cli_cmds.py,bash -n scripts/task_test_jit_cache_package_build_import.sh,git diff --check, workflow YAML parsing, and the DGX Spark CLI dry-run.SM110 architecture split:
11.0a/sm110jit-cache build coverage to CUDA 13.0aarch64release, nightly, and PR AOT build/import arch lists. CUDA 13.0x86_64lists now omit11.0a.11.0afrom the generic x86-oriented examples and call out adding it for Jetson AGX Thor / T5000 aarch64 targets.bash -n scripts/task_test_jit_cache_package_build_import.sh,git diff --check,uvx ruff check tests/cli/test_cli_cmds.py, andpython -m pytest tests/cli/test_cli_cmds.py -q(27 passed).x86_64->7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0f;aarch64->7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f.Summary by CodeRabbit
New Features
Documentation
Tests / CI
Chores