Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughFP4 quantization module dispatch now validates the backend and conditionally selects the AOT ChangesFP4 backend dispatch
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to The change selects architecture-specific FP4 modules for SM12x while preserving prebuilt wheel fallback behavior; no actionable merge-blocking risk remains after normal checks and review. Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 the backend selection logic for SM12x architectures in fp4_quantization.py to prefer arch-specific modules over the family variant, falling back to the family variant only when it is prebuilt (AOT) and the arch-specific one is not. The reviewer suggested optimizing this logic by instantiating the module generators once to avoid redundant calls and return early.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/quantization/fp4_quantization.py (1)
237-248: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider using
backend_modules["120f"]()instead ofgen_fp4_quantization_sm120f_module()for consistency.Line 247 uses the direct function reference
gen_fp4_quantization_sm120f_module()while line 246 uses the dictionary lookupbackend_modules[backend](). Since"120f"is already a key inbackend_modules, usingbackend_modules["120f"]()would be more consistent and make the relationship between the two checks clearer.♻️ Consistency tweak
if ( version_at_least(torch.version.cuda, "12.9") and not backend_modules[backend]().is_aot - and gen_fp4_quantization_sm120f_module().is_aot + and backend_modules["120f"]().is_aot ):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/quantization/fp4_quantization.py` around lines 237 - 248, Replace the direct gen_fp4_quantization_sm120f_module() call in the SM12x dispatch condition with backend_modules["120f"](), keeping the existing AOT check and aligning it with the backend_modules[backend]() lookup.
🤖 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.
Nitpick comments:
In `@flashinfer/quantization/fp4_quantization.py`:
- Around line 237-248: Replace the direct gen_fp4_quantization_sm120f_module()
call in the SM12x dispatch condition with backend_modules["120f"](), keeping the
existing AOT check and aligning it with the backend_modules[backend]() lookup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5945f2e3-434f-455f-8dba-595100daa324
📒 Files selected for processing (1)
flashinfer/quantization/fp4_quantization.py
|
@coderabbitai The |
|
Tip For best results, initiate chat on the files or code changes.
ᕦ(ò_óˇ)ᕤ |
|
Verified on an RTX PRO 6000 (SM120, CUDA 13.0, JIT path, no AOT dir) at 3574fd1: the runtime now builds and loads fp4_quantization_120 with -gencode=arch=compute_120a,code=sm_120a where main picks the 120f module on the same box, and the full tests/utils/test_fp4_quantize.py passes 10306/10306 vs main's 10300/10300 baseline (the +6 are this PR's selection tests). The AOT-redirect leg (release wheel shipping only 120f) is the one path I couldn't exercise here -- no AOT directory on this box. |
get_fp4_quantization_module() unconditionally redirected backends "120" and "121" to "120f" on CUDA >= 12.9, while aot.py builds the arch-specific fp4_quantization_120/_121 modules for 12.0a/12.1a arch lists. The AOT artifacts were dead weight, and JIT-cache builds without 12.0f in the arch list broke FLASHINFER_DISABLE_JIT (or silently recompiled) because the runtime asked for a module nobody built. Use the arch-specific module (sm_120a / sm_121a, strict supersets of sm_120f) directly, matching GEMM/MoE dispatch. Redirect to 120f only when it is the module available ahead-of-time, so release wheels (built with 12.0f only) keep loading their prebuilt module. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… test Address review feedback on the arch-specific dispatch: - Bind the JitSpec once instead of constructing it for the AOT probe and again for build_and_load. - Drop the CUDA >= 12.9 gate: the redirect loads a prebuilt .so, so the toolkit version is irrelevant, and the non-AOT path keeps the same arch flags either way. - Reword the comment; the old text wrongly claimed release wheels prebuild only sm_120f (a 12.0f wheel also prebuilds sm_120a via aot.py's substring match, so only backend "121" ever redirects) and overstated the parallel with GEMM/MoE dispatch. - Add a CPU-only unit test for the selection logic, which CI never exercises through a real AOT directory. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The class-level is_aot monkeypatch broke when JitSpec grew subclasses (the subclass property shadows the base-class patch). Point FLASHINFER_AOT_DIR at a tmp dir with placeholder .so files so the production is_aot check runs unmodified. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3574fd1 to
abd7b0c
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
3 participants
📌 Description
Resolves item 15 of the SM121 support audit (#3170): the FP4 quantization runtime and AOT disagree about which module to use on SM12x.
get_fp4_quantization_module()unconditionally redirected backends"120"/"121"to"120f"on CUDA >= 12.9 (introduced in #2650), whileaot.pybuilds the arch-specificfp4_quantization_120/_121modules. The AOT artifacts for12.0a/12.1awere dead weight, and a JIT-cache build without12.0fin the arch list brokeFLASHINFER_DISABLE_JIT.Fixes:
sm_120a/sm_121a) directly.120fonly when it is the module available ahead-of-time, so release wheels keep loading their prebuilt module with no new JIT compile.🔍 Related Issues
#3170 (item 15), #2650 (introduced the redirect).
🧪 Tests
On DGX Spark (GB10, SM121, CUDA 13, source install):
test_fp4_quantization_module_selectioncovering the arch-vs-family selection.tests/utils/test_fp4_quantize.py(10300 passed) andtest_fp4_quantize_padding.py(16 passed); the JIT cache now buildsfp4_quantization_121withcompute_121a.sm_120fmodule staged inFLASHINFER_AOT_DIR, a fresh process loads it and round-trips correctly.Reviewer Notes
120fJIT cache no longer matches).sm_120fwhile source installs compilesm_121a, matching the existing SM12x GEMM/MoE behavior..so, so the toolkit version is irrelevant.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests