Conversation
CUDA >= 12.9 flashinfer-jit-cache wheels build only the 12.0f family target for SM12x, so aot.py never detects sm121 and skips the arch-specific fp4_quantization_121 module; generic modules also carry no sm_121a SASS. Add 12.1a alongside 12.0f in the aarch64 lists for CUDA 12.9 and 13.0 — SM121 (DGX Spark) hardware is aarch64-only, following the 11.0a (Thor) precedent — and keep x86_64 lists unchanged. Sync the mirror list in the jit-cache package build test script and the documented example arch lists. AI-assisted with Claude Code. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The flashinfer-jit-cache example FLASHINFER_CUDA_ARCH_LIST mixes archs with different minimum toolkit versions (11.0a needs CUDA 13.0, 12.0f and 12.1a need 12.9), so copying it verbatim fails on older toolkits. Add a short note next to each documented list. AI-assisted with Claude Code. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR aligns JIT cache CUDA architecture selection across build scripts and release workflows, adds aarch64-specific targets, and updates documentation with CUDA 13.0 compatibility guidance. ChangesJIT cache CUDA architecture selection
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![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 documentation and build scripts to support the 12.1a CUDA architecture on aarch64 platforms for CUDA versions 12.9 and above. Specifically, it updates the FLASHINFER_CUDA_ARCH_LIST examples in CLAUDE.md, README.md, and docs/installation.rst, and updates the JIT cache package build script to conditionally append 12.1a based on the CUDA version and architecture. Feedback was provided to consolidate redundant conditional blocks in the JIT cache build script to improve maintainability.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if (major, minor) >= (13, 0): | ||
| arches.append("10.0a") | ||
| arches.append("10.3a") | ||
| machine = (os.environ.get("ARCH") or platform.machine()).lower() | ||
| if machine in ("aarch64", "arm64"): | ||
| if is_aarch64: | ||
| arches.append("11.0a") | ||
| arches.append("12.0f") | ||
| if is_aarch64: | ||
| arches.append("12.1a") | ||
| elif (major, minor) >= (12, 9): | ||
| arches.append("10.0a") | ||
| arches.append("10.3a") | ||
| arches.append("12.0f") | ||
| if is_aarch64: | ||
| arches.append("12.1a") |
There was a problem hiding this comment.
The logic blocks for CUDA >= 13.0 and >= 12.9 are highly redundant, differing only by the conditional addition of 11.0a on aarch64 for CUDA >= 13.0. Consolidating these blocks reduces duplication and improves maintainability while preserving the exact order of the architectures.
| if (major, minor) >= (13, 0): | |
| arches.append("10.0a") | |
| arches.append("10.3a") | |
| machine = (os.environ.get("ARCH") or platform.machine()).lower() | |
| if machine in ("aarch64", "arm64"): | |
| if is_aarch64: | |
| arches.append("11.0a") | |
| arches.append("12.0f") | |
| if is_aarch64: | |
| arches.append("12.1a") | |
| elif (major, minor) >= (12, 9): | |
| arches.append("10.0a") | |
| arches.append("10.3a") | |
| arches.append("12.0f") | |
| if is_aarch64: | |
| arches.append("12.1a") | |
| if (major, minor) >= (12, 9): | |
| arches.append("10.0a") | |
| arches.append("10.3a") | |
| if (major, minor) >= (13, 0) and is_aarch64: | |
| arches.append("11.0a") | |
| arches.append("12.0f") | |
| if is_aarch64: | |
| arches.append("12.1a") |
There was a problem hiding this comment.
Thanks — the suggestion is output-equivalent, but keeping the branches separate is intentional: the three-branch if/elif maps one-to-one onto the workflow ternary's branches in release.yml/nightly-release.yml (12.8 / 12.9 / 13.0), so each branch reads as that CUDA version's complete list and can be eyeball-diffed against the workflow string directly. Merging 12.9 and 13.0 would save three lines but lose that mapping, so we'll keep the current structure.
|
@flashinfer-bot run |
|
Size measurement for reviewers (method: per-module readelf/cuobjdump analysis of the published v0.6.9 cu130 aarch64 jit-cache wheel, calibrated against a local build with and without the extra gencode):
|
|
It looks like right now, the cu129 nightly builds are still large enough that adding 121a will cause them to spill over the 2GiB limit significantly: https://github.com/flashinfer-ai/flashinfer/releases/tag/nightly-v0.6.15-20260722 I'm really not sure how to proceed at this point other than getting back to trying to split up the jit-cache wheels 😞 |
Good point! How about I scope this PR down to cu130 only? The only sm_121 device is GB10, which ships with CUDA 13 anyway, and cu129 already covers it via 12.0f. And if this feels too messy, I'm also fine shelving this PR until the wheel split is implemented. |
4 participants
📌 Description
Resolves item 12 of the SM121 support audit (#3170). The CUDA >= 12.9
flashinfer-jit-cachewheels build only the12.0ffamily target for SM12x. Twoconsequences:
detect_sm_capabilities()inaot.pylooks forcompute_121in the gencode flags,never finds it, and skips the arch-specific
fp4_quantization_121module. The SM120counterpart is prebuilt anyway, because
has_sm120matches thecompute_120substringinside
compute_120f.sm_121aSASS. The family target covers themfunctionally today; an arch build leaves room for arch-specific features later (for
example sparse block-scaled MMA).
Fixes:
12.1anext to12.0fin the aarch64FLASHINFER_CUDA_ARCH_LISTfor the CUDA12.9 and 13.0 wheel builds (
release.yml,nightly-release.yml).scripts/task_test_jit_cache_package_build_import.sh.docs/installation.rst, andCLAUDE.md,which also gains the missing
10.3a), and add a short note that the full list needsCUDA 13.0 and should be trimmed on older toolkits (
11.0aalready required 13.0before this change).
and it already ships
12.0a.🔍 Related Issues
#3170 (item 12). Pairs with #3906 (item 15): together they let SM121 wheel installs load
the native
sm_121aFP4 quantization module instead of falling back to120f, matchingsource installs and keeping
FLASHINFER_DISABLE_JITcovered.🧪 Tests
On DGX Spark (GB10, SM121, CUDA 13, aarch64):
numbers for
<; for these values that matches plain string comparison) and the scriptsnippet for all 6 matrix combos
({12.8, 12.9, 13.0} x {x86_64, aarch64}). Both produce identical lists that match the
intended table, and only the two aarch64 >= 12.9 lists change.
detect_sm_capabilities()reportssm121: True, andCompilationContextemits thecompute_120fandcompute_121agencodes together.normmodule with both gencodes in one list;cuobjdump --list-elfshows thesm_120(family) andsm_121acubins side by side in one fatbin.fp4_quantization_121builds withcompute_121a.pre-commit runon the changed files: clean.Reviewer Notes
12.0a; skipped on purpose. The item's reasoning ("wheels lackNVFP4/MXFP4 MMA") was corrected later in the issue: dense block-scaled FP4 MMA, the only
kind FlashInfer uses, is available on
120f(PTX ISA >= 8.8), and the SM120 arch FP4module is already prebuilt through the substring match above. Adding
12.0awould onlygrow the wheels.
12.1agoes into the aarch64 lists only, like11.0a(Thor), since SM121 hardware hasno x86 host.
major version 12, so larger fatbins and longer CI builds. x86_64 wheels are unchanged.
🤖 Generated with Claude Code
Summary by CodeRabbit
Improvements
Documentation
flashinfer-jit-cacheto note that the full CUDA architecture list requires CUDA 13.0 and should be trimmed for older toolkits.FLASHINFER_CUDA_ARCH_LISTexamples to include the newly supported architecture entries.