fix(jit): warn when CUDA toolkit < 12.8 drops capability flags from JIT builds - #4429
Open
rishabhsinha17 wants to merge 2 commits into
Open
rishabhsinha17 wants to merge 2 commits into
rishabhsinha17 wants to merge 2 commits into
Conversation
rishabhsinha17
requested review from
Aneureka,
IwakuraRein,
StudyingShao,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
August 9, 2026 17:12
Contributor
|
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 (3)
Included review availability: Your plan includes up to 8 reviews per rolling hour; 7 remain after this review. 📝 WalkthroughWalkthroughThe PR adds shared CUDA toolkit and AOT-module gating for JIT compiler flags. FP4 and FP8 paths use the gating logic. Fused-MoE validation reports unsupported builds and toolkit versions. CPU-only tests cover toolkit and AOT scenarios. ChangesToolkit-gated JIT support
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to The change makes unsupported CUDA toolkit behavior explicit and preserves AOT-backed execution paths; no actionable merge-blocking risk remains after normal checks and review. Sequence Diagram(s)sequenceDiagram
participant Caller
participant fused_moe
participant cpp_ext
participant AOTCache
Caller->>fused_moe: request FP8 block-scale capability
fused_moe->>cpp_ext: check CUDA version and AOT module
cpp_ext->>AOTCache: inspect matching prebuilt module
AOTCache-->>cpp_ext: module present or absent
cpp_ext-->>fused_moe: capability result
fused_moe-->>Caller: supported or unsupported
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![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: 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/cpp_ext.py`:
- Around line 94-101: The AOT presence check must not imply that the artifact is
usable. Update has_prebuilt_aot_module() and the related support-probe gating to
validate the AOT load result, or track presence separately from successful
loading, so stale, corrupt, or incompatible artifacts fall back through
JitSpecNvcc.try_load() and do not report unsupported CUDA paths as supported;
retain the JIT-flag warning downgrade only after usability is confirmed. Update
tests/jit/test_toolkit_version_gated_flags.py lines 44-48 to cover an unloadable
AOT artifact and verify the support probe and warning behavior.
🪄 Autofix
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 Plus
Run ID: 757b9371-69d8-48f8-9396-d0c40616f572
📥 Commits
Reviewing files that changed from the base of the PR and between 4fbac49 and 069e17a3db6d2ccde91fc28dd38ce4a3b8cf69c8.
📒 Files selected for processing (7)
flashinfer/fused_moe/core.pyflashinfer/jit/cpp_ext.pyflashinfer/jit/fp4_quantization.pyflashinfer/jit/fused_moe.pyflashinfer/jit/gemm/fp8_blockscale.pyflashinfer/quantization/fp4_quantization.pytests/jit/test_toolkit_version_gated_flags.py
…IT builds Flags gated on the JIT toolkit version (-DENABLE_FP4 in fp4_quantization, -DENABLE_FP8_BLOCK_SCALE in cutlass fused_moe 89/90 and fp8_blockscale_gemm_90) were dropped silently, so the module built fine and the first gated call failed deep in kernel-runner construction with a misleading dtype-combination error (flashinfer-ai#3951). - Add version_gated_nvcc_flag() which emits a loud warning naming the module, the dropped flag, the toolkit version found, and the 12.8 requirement. - Use it at every version-gated flag drop site: jit/fused_moe.py, jit/fp4_quantization.py, quantization/fp4_quantization.py, and jit/gemm/fp8_blockscale.py. - Demote the warning to info when a prebuilt AOT artifact exists for the module: JitSpecNvcc.try_load() serves such modules from the AOT cache, so the JIT flag list is never compiled and the gated kernels stay available. Without this, the recommended flashinfer-jit-cache install with an older local toolkit warned spuriously, even recommending the wheel already installed. The check mirrors the AOT resolution rule through a new has_prebuilt_aot_module() helper (shared with the capability probe below) instead of moving the warning to build time, which would require threading dropped-flag metadata through JitSpec for a strictly larger diff. AOT wheel builds themselves compile into a separate build dir and never read FLASHINFER_AOT_DIR, so self-builds with an old toolkit still warn loudly. - Make the existing use_deepseek_fp8_block_scale early check name the actual cause and stop rejecting AOT/jit-cache installs, via a new cutlass_fused_moe_fp8_block_scale_supported() capability probe. - Document why SM90 fused_moe keeps -DENABLE_FP4 unconditionally (cutlass::float_e2m1_t path from flashinfer-ai#3738, no <cuda_fp4.h> needed). - Add CPU-only tests that monkeypatch the toolkit-version probe to both sides of 12.8 and assert the flag, warning, and AOT-demotion behavior; tests isolate FLASHINFER_AOT_DIR so host jit-cache installs cannot skew results.
rishabhsinha17
force-pushed
the
fix/fused-moe-fp4-toolkit-guard
branch
from
August 9, 2026 17:21
069e17a to
9ac29d2
Compare
Author
Issue flashinfer-ai#3951 request 2: frameworks selecting between MoE backends need a public probe for fp4 availability before committing to a multi-minute JIT build; the only options were probing private jit internals or try/except. - Add has_fp4_support(device_arch="100") to jit/fp4_quantization.py, mirroring cutlass_fused_moe_fp8_block_scale_supported(): True when a prebuilt AOT fp4_quantization module for the arch exists or the local CUDA toolkit is 12.8+, i.e. when -DENABLE_FP4 survives version_gated_nvcc_flag(). Presence-based, no module load. - Re-export from flashinfer/jit/__init__.py so the spelling requested in the issue, flashinfer.jit.has_fp4_support(), works. - Extend the CPU-only gated-flags tests: toolkit 12.4 -> False, 12.8 -> True (no-arg form), AOT artifact with a 12.4 toolkit -> True for the covered arch only, and the flashinfer.jit export identity.
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
When the JIT toolchain is older than CUDA 12.8, version-gated nvcc flags are dropped silently:
-DENABLE_FP4in the fp4_quantization generators and-DENABLE_FP8_BLOCK_SCALEin the CUTLASS fused_moe 89/90 and fp8_blockscale_gemm_90 generators. The module builds fine without the capability and the first gated call fails deep in kernel-runner construction with a misleading dtype-combination error (see #3951, hit through vLLM W4A16-MXFP4 serving on a 12.4 toolkit).version_gated_nvcc_flag()inflashinfer/jit/cpp_ext.py: dropping a flag now logs a warning at flag-generation time naming the module, the flag, the toolkit version found, and the 12.8 requirement.jit/fused_moe.py,jit/fp4_quantization.py,quantization/fp4_quantization.py,jit/gemm/fp8_blockscale.py). One warning per generated module, not one per flag list.try_load()serves it and the JIT flag list is never compiled, so the drop is logged at info level instead of warning. The check goes through a newhas_prebuilt_aot_module()helper mirroringJitSpecNvcc.aot_path. AOT wheel builds compile into a separate build dir and never readFLASHINFER_AOT_DIR, so self-builds with an old toolkit still warn loudly.use_deepseek_fp8_block_scaleearly error ("not implemented for CUDA 12.6 or lower") infused_moe/core.pywith one naming the actual cause and remedies, backed by a newcutlass_fused_moe_fp8_block_scale_supported()capability probe. The probe also accepts a prebuilt AOTfused_moe_90module, so flashinfer-jit-cache installs with an older local toolkit are no longer falsely rejected.-DENABLE_FP4unconditionally after perf: optimize MXFP4xBF16 & INT4xFP8 and add MXFP4xFP8 CUTLASS MoE backend for SM90 #3738 (Hopper path usescutlass::float_e2m1_tfromfp4_compat.h, no<cuda_fp4.h>needed) and pin it with a regression test, since that unconditional flag is what fixed the original fp4 repro on main.🔍 Related Issues
Closes #3951
🚀 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.).tests/jit/test_toolkit_version_gated_flags.py(14 tests) is CPU-only: it monkeypatches the toolkit-version probe to 12.4 and 12.8 and asserts flag presence, warning content (module name, dropped flag, found and required versions), warning count, info-level demotion when an AOT artifact exists, and the AOT escape hatch of the capability probe. Tests isolateFLASHINFER_AOT_DIRso a host jit-cache install cannot skew results. No nvcc or GPU required.Reviewer Notes
Two behavior changes with flashinfer-jit-cache installed and a local toolkit < 12.8:
cutlass_fused_moe(use_deepseek_fp8_block_scale=True)previously raised NotImplementedError based on the local toolkit alone and now proceeds because the AOT module contains the kernels, and flag generation for AOT-served modules logs info instead of warning. Both assume AOT modules are built with a 12.8+ toolkit, which holds for the official wheels (cu128/cu129/cu130). Issue request 2 (a publichas_fp4_support()query) is left out to keep this focused; the capability-probe pattern here is the template for it.Summary by CodeRabbit
New Features
Bug Fixes
Tests