perf(jit): drop dead SM12x gencode from SM10x-serving modules by yichengj0 · Pull Request #3947 · flashinfer-ai/flashinfer

yichengj0 · 2026-07-13T17:59:30Z

📌 Description

Resolves cross-cutting item 10 of the SM121 support audit (#3170): several SM10x-serving
JIT module generators include major version 12 in supported_major_versions, so wheel
builds whose arch list contains a 12.x entry compile their kernels for SM12x too. That
SM12x machine code can never run: on a compute capability 12.x device, dispatch always
selects the dedicated sm120 module (cutlass fused-moe, groupwise gemm, cutlass
fp4/fp8/bf16 gemm) or the sm100 variant (mamba selective_state_update), and on SM10x
devices the CUDA loader picks the matching sm_100a/sm_103a code from the fatbin.

Measured in the published v0.6.9 cu130 aarch64 jit-cache wheel (per-module
readelf/cuobjdump accounting), this dead code is 1,046 MB installed, 52% of all SM12x
bytes in the wheel:

module	dead SM12x bytes
`fused_moe_103`	430 MB
`fused_moe_100`	409 MB
`gemm_sm100`	48 MB
`fp8_gemm_cutlass`	8 MB
`fp4_gemm_cutlass`	5 MB
`selective_state_update` base/sm90 variants (100 modules)	~147 MB

Fixes:

Remove 12 from supported_major_versions at the nine dead sites:
gen_cutlass_fused_moe_sm100_module, gen_cutlass_fused_moe_sm103_module
(flashinfer/jit/fused_moe.py); gen_gemm_sm100_module,
gen_gemm_sm100_module_cutlass_fp4, gen_gemm_sm103_module_cutlass_fp4,
gen_gemm_sm100_module_cutlass_fp8, gen_gemm_sm100_module_cutlass_bf16
(flashinfer/jit/gemm/core.py); gen_selective_state_update_module,
gen_selective_state_update_sm90_module (flashinfer/jit/mamba/selective_state_update.py).

This matters for the release pipeline right now: the cu130 wheels sit just under GitHub's
2 GiB release-asset cap (aarch64 at 2.117 GB, x86_64 at 2.133 GB, cap 2.147 GB), so any
module growth breaks the next upload. This change shrinks the aarch64 wheel to roughly
1.88 GB zipped (x86_64 similar).

🔍 Related Issues

#3170 (cross-cutting item 10).

🧪 Tests

A dispatch audit shows that CC 12.x never reaches the nine changed generator sites.
The change was verified at runtime on both SM121 (DGX Spark, GB10) and SM120
(RTX 5080). Starting from an empty JIT cache, the test
suites for all six affected op families (mamba selective state update, bmm_fp8,
mm_fp4, mm_bf16, groupwise-scaled fp8 GEMM, and cutlass fused MoE) pass on both
machines, and the cache afterwards contains none of the nine modules this PR edits.
This confirms that SM12x dispatch never loads them.
pre-commit run on the changed files is clean.

Reviewer Notes

The dead verdicts depend on today's dispatch gates (get_cutlass_fused_moe_module's
backend strings, is_sm12x_supported branches preceding is_sm100a_supported, the
bmm_fp8 heuristic preferring cutlass_sm12x, and the mamba _get_module ordering).
If a future change routes SM12x to one of these modules, its generator needs 12 back.
The mamba base/sm90 variants also embed SM9x/SM10x code that the same dispatch logic
never loads (base only serves SM8x, sm90 only SM9x). Not touched here to keep this PR
scoped to the SM12x audit; a follow-up could prune those the same way.
No runtime behavior changes on any device: no dispatch path is modified, only the set
of arch targets each module is compiled for.

🤖 Generated with Claude Code

Summary by CodeRabbit

Bug Fixes
- Improved CUDA architecture compatibility for fused MoE, GEMM, and selective state update operations.
- Prevented unsupported CUDA 12 compilation targets for selected GPU kernels, while preserving CUDA 12 support where applicable.

coderabbitai · 2026-07-13T17:59:55Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 960983e3-f242-4421-beb7-bc0144e4a2cc

📥 Commits

Reviewing files that changed from the base of the PR and between 190aee650be2968f5807d59845f2f4cecacf3ad1 and 4053636.

📒 Files selected for processing (3)

flashinfer/jit/fused_moe.py
flashinfer/jit/gemm/core.py
flashinfer/jit/mamba/selective_state_update.py

🚧 Files skipped from review as they are similar to previous changes (2)

flashinfer/jit/mamba/selective_state_update.py
flashinfer/jit/fused_moe.py

📝 Walkthrough

Walkthrough

NVCC architecture flag selection is narrowed for fused MoE, GEMM, and selective state update JIT generators by removing CUDA major version 12 from non-SM120 compilation paths.

Changes

CUDA architecture dispatch

Layer / File(s)	Summary
Fused MoE dispatch `flashinfer/jit/fused_moe.py`	SM103 supports CUDA major version 10, while SM100 supports versions 10 and 11.
GEMM dispatch `flashinfer/jit/gemm/core.py`	Five SM100/SM103 GEMM generators change supported CUDA major versions from `[10, 11, 12]` to `[10, 11]`.
Selective state update dispatch `flashinfer/jit/mamba/selective_state_update.py`	Generic and SM90 generators remove CUDA major version 12 from their supported versions.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

flashinfer-ai/flashinfer#2560: Narrows JIT NVCC targeting by removing CUDA major version 12.
flashinfer-ai/flashinfer#2725: Modifies fused MoE CUDA major-version dispatch at overlapping code points.
flashinfer-ai/flashinfer#2082: Modifies SM103 fused MoE CUDA-architecture JIT generation.

Suggested reviewers: saltyminty, yzh119, bkryu

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: removing dead SM12x gencode from SM10x-serving JIT modules.
Description check	✅ Passed	The description largely follows the template and includes Description, Related Issues, Tests, and Reviewer Notes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

gemini-code-assist

Code Review

This pull request removes major version 12 from the supported_major_versions list in several JIT module generation functions across fused_moe.py, gemm/core.py, and selective_state_update.py. This change prevents redundant compilation, as SM12x dispatches to other modules (such as sm120 or sm100 variants). Explanatory comments have been added to document this dispatch behavior. I have no further feedback to provide as the changes are clean and well-documented.

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.

aleozlx · 2026-07-15T04:10:39Z

@@ -226,7 +226,7 @@ def gen_selective_state_update_sm90_module(
    )
    compilation_context = CompilationContext()
    nvcc_flags = compilation_context.get_nvcc_flags_list(
-        supported_major_versions=[9, 10, 11, 12]
+        supported_major_versions=[9, 10, 11]


so should we drop 10, 11 as well based on the same pattern that this is sm90 module?

actually let make my question about just "10", because i remember 11 may be piggy backed on to some existing code path, risky to drop without auditing how 11 works.

confirmed dead. will fix this in a follow up pr.

aleozlx · 2026-07-15T04:14:43Z

@@ -173,7 +173,7 @@ def gen_selective_state_update_module(
    # "No supported CUDA architectures found" instead of failing in nvcc.
    compilation_context = CompilationContext()
    nvcc_flags = compilation_context.get_nvcc_flags_list(
-        supported_major_versions=[8, 9, 10, 11, 12]
+        supported_major_versions=[8, 9, 10, 11]


should we drop 9 from here, given sm90 is below?

aleozlx · 2026-07-15T04:16:32Z

/bot run

flashinfer-bot · 2026-07-15T04:17:27Z

GitLab MR !973 has been created, and the CI pipeline #58066888 is currently running. I'll report back once the pipeline job completes.

Nine JIT module generators included major version 12 in their supported_major_versions even though runtime dispatch on SM12x devices always selects the dedicated sm120 (cutlass fused-moe, groupwise gemm, cutlass fp4/fp8/bf16 gemm) or sm100 (mamba selective_state_update) variants, so the SM12x machine code they embed is never loaded. In the cu130 aarch64 jit-cache wheel this dead code is 1,046 MB installed, 52% of all SM12x bytes: fused_moe_103 (430 MB), fused_moe_100 (409 MB), gemm_sm100 (48 MB), fp8_gemm_cutlass (8 MB), fp4_gemm_cutlass (5 MB), and ~147 MB across the selective_state_update base/sm90 variants. Remove 12 from those generators. A dispatch audit confirmed each one is unreachable from compute capability 12.x, that generators genuinely shared with SM12x (trtllm fused-moe, bgmv_moe, tinygemm2, cublaslt bf16 gemm, xqa, dcp_alltoall, fp4 kv quant/dequant, fa2 large-head-dim, fmha_v2) keep 12, and that nothing invokes the changed generators on a machine whose arch list contains only 12.x. AI-assisted with Claude Code. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

flashinfer-bot · 2026-07-15T09:25:08Z

[FAILED] Pipeline #58066888: 4/20 passed

aleozlx · 2026-07-17T05:44:30Z

/bot run tests/gemm

flashinfer-bot · 2026-07-17T05:45:32Z

GitLab MR !973 has been updated with latest changes, and the CI pipeline #58402720 is currently running. I'll report back once the pipeline job completes.

bkryu · 2026-07-17T16:39:37Z

/bot run tests/mamba

flashinfer-bot · 2026-07-17T16:40:23Z

GitLab MR !973 has been created, and the CI pipeline #58467724 is currently running. I'll report back once the pipeline job completes.

bkryu · 2026-07-17T20:09:21Z

bot run /tests/moe

bkryu · 2026-07-17T20:34:07Z

/bot run tests/moe

flashinfer-bot · 2026-07-17T20:34:34Z

GitLab MR !973 has been created, and the CI pipeline #58496533 is currently running. I'll report back once the pipeline job completes.

kahyunnam

lgtm

flashinfer-bot · 2026-07-17T22:57:44Z

[FAILED] Pipeline #58496533: 10/20 passed

## 📌 Description Follow-up to #3947, which removed compiled GPU code that SM12x devices can never run. This PR extends the same audit to the whole tree: every module's `supported_major_versions` list was checked against the dispatch logic that decides which module a GPU actually loads. Four lists still contained architectures whose GPUs can never load the module: | module | change | why the removed entries are dead | |---|---|---| | mamba SSU base | `[8, 9, 10, 11]` → `[8]` | SM90 and newer GPUs always load the sm90 or sm100 variant instead | | mamba SSU sm90 | `[9, 10, 11]` → `[9]` | SM100 and newer GPUs always load the sm100 variant | | fp8 GEMM (`gemm_sm100`) | `[10, 11]` → `[10]` | every caller accepts SM100 only; SM110 gets an error before the module loads | | xqa | `[9, 10, 11, 12]` → `[9, 10, 12]` | the xqa API rejects SM110 before the module can load | `aot.py` gets a matching change: the two mamba modules are now prebuilt only for the build targets that load them (base for SM8x, sm90 variant for SM9x). This is required: after the prune, prebuilding the base module in a wheel with no SM8x target has no architecture left to compile for and fails the build. ### Result This further removes about 587 MB of installed size from the cu130 aarch64 jit-cache wheel, on top of the ~1 GB removed by #3947: | pruned GPU code | installed | |---|---| | mamba SSU base variants (48 modules) | 214 MB | | mamba SSU sm90 variants (48 modules) | 267 MB | | xqa (384 modules), SM110 code | 56 MB | | fp8 GEMM (`gemm_sm100`), SM110 code | 50 MB | The xqa and fp8 GEMM rows are SM110 code, which only exists in cu130 aarch64 wheels (the only build that targets SM110). The mamba rows apply to every wheel. ## 🔍 Related Issues #3170 (cross-cutting item 10, extended beyond SM12x). ## 🧪 Tests - A GPU-free dispatch simulation (45 cases): for every compute capability from 8.0 to 12.1, dispatch picks the expected module, and that module still compiles for the architecture. All pass on this branch. Removing a live entry on purpose makes exactly the matching case fail, which confirms the check can catch a bad prune. - Runtime verification on H100 (SM90) and B200 (SM100), JIT-compiling this branch into a fresh cache. All affected suites pass: mamba SSU and xqa on both machines, and the `gemm_sm100` fp8 GEMM suite on B200. The cache afterwards holds only sm90 variants on H100 and only sm100 variants on B200, so dispatch never requested the pruned modules. - `pre-commit run` on the changed files is clean. ## Reviewer Notes - Two more dead entries (SM110 in the sm103 fp4 and cutlass bf16 GEMM modules) stay in place: those modules are not prebuilt into wheels, so keeping the entries costs nothing and future SM110 support there needs no list edit. - The dead verdicts reflect today's dispatch. If a later change routes one of these architectures back (for example, Thor support in xqa), the entry must be re-added here. - The audit also found the opposite mismatch: some modules do run on architectures that wheels never prebuild them for (several SM110-capable modules, and the trtllm-gen fused MoE on SM12x). Left unchanged, since adding prebuilds grows wheels and deserves its own decision. 🤖 Generated with [Claude Code](https://claude.com/claude-code)  ## Summary by CodeRabbit * **Bug Fixes** * Improved GPU architecture targeting for generated CUDA kernels. * Corrected architecture-specific generation for selective state update, GEMM, and XQA kernels. * Reduced compilation attempts for unsupported GPU architectures, improving build compatibility and reliability.  Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…ache wheel (#4527) ## 📌 Description #3684 added `12.1a` to the aarch64 arch lists for **both** cu129 and cu130. Wheel sizes from the Actions artifact API, last good nightly (Aug 11) vs Aug 14: | wheel | Aug 11 | Aug 14 | headroom vs 2 GiB | | --- | --- | --- | --- | | cu128 x86_64 | 1.238 | 1.274 | 743 MiB | | cu128 aarch64 | 1.228 | 1.264 | 753 MiB | | cu129 x86_64 | 1.861 | 1.906 | 96 MiB | | **cu129 aarch64** | 1.848 | **2.143** | **−147 MiB** | | cu130 x86_64 | 1.445 | 1.480 | 533 MiB | | cu130 aarch64 | 1.619 | 1.852 | 151 MiB | Only cu129 aarch64 is actually over. The cu130 wheels are missing from those releases as collateral: the upload loop is `for cuda in 128 129 130; do for arch in x86_64 aarch64` under `shell: bash -e`, so it aborts on the 4th asset and never attempts the cu130 pair. Subtracting each aarch64 wheel's delta from its same-CUDA x86_64 delta prices the new target at **+256 MiB** on cu129 aarch64 and +204 MiB on cu130 aarch64. Removing it from cu129 aarch64 lands that wheel at roughly **1.893 GiB** (~110 MiB headroom). cu130 aarch64 keeps `12.1a` and fits at 1.852 GiB. ## 🔍 Related Issues Fixes #4519. Caused by #3684. Wheel splitting (#4514) remains the durable fix — cu129 x86_64 has only 96 MiB left and grew ~2.8 MiB/day between Jul 22 and Aug 12. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). Workflow-only change; there is no unit test for the release matrix. Verified by parsing both workflows as YAML and evaluating the expression for all six matrix entries (table above). ## Reviewer Notes Two questions worth a maintainer's call before this comes out of draft: 1. @jethachan — does SM121 need a warm AOT cache on **CUDA 12.9 aarch64** specifically, or is cu130 aarch64 sufficient? #3684's validation notes cover GB10 and RTX PRO 6000 but do not say which CUDA minor was used. 2. The alternative is to keep `12.1a` and instead land per-module filtering in the style of #3947. That needs a minor-version filter added to `get_nvcc_flags_list`, which today filters on major version only — more invasive than seems wise for v0.6.18.  ## Summary by CodeRabbit * **Bug Fixes** * Updated CUDA 12.9 wheel builds to use a consistent architecture set across platforms. * Preserved the specialized architecture target for CUDA 13.0 ARM64 builds. * Documented wheel-size constraints for affected builds.

…ache wheel (flashinfer-ai#4527) ## 📌 Description flashinfer-ai#3684 added `12.1a` to the aarch64 arch lists for **both** cu129 and cu130. Wheel sizes from the Actions artifact API, last good nightly (Aug 11) vs Aug 14: | wheel | Aug 11 | Aug 14 | headroom vs 2 GiB | | --- | --- | --- | --- | | cu128 x86_64 | 1.238 | 1.274 | 743 MiB | | cu128 aarch64 | 1.228 | 1.264 | 753 MiB | | cu129 x86_64 | 1.861 | 1.906 | 96 MiB | | **cu129 aarch64** | 1.848 | **2.143** | **−147 MiB** | | cu130 x86_64 | 1.445 | 1.480 | 533 MiB | | cu130 aarch64 | 1.619 | 1.852 | 151 MiB | Only cu129 aarch64 is actually over. The cu130 wheels are missing from those releases as collateral: the upload loop is `for cuda in 128 129 130; do for arch in x86_64 aarch64` under `shell: bash -e`, so it aborts on the 4th asset and never attempts the cu130 pair. Subtracting each aarch64 wheel's delta from its same-CUDA x86_64 delta prices the new target at **+256 MiB** on cu129 aarch64 and +204 MiB on cu130 aarch64. Removing it from cu129 aarch64 lands that wheel at roughly **1.893 GiB** (~110 MiB headroom). cu130 aarch64 keeps `12.1a` and fits at 1.852 GiB. ## 🔍 Related Issues Fixes flashinfer-ai#4519. Caused by flashinfer-ai#3684. Wheel splitting (flashinfer-ai#4514) remains the durable fix — cu129 x86_64 has only 96 MiB left and grew ~2.8 MiB/day between Jul 22 and Aug 12. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). Workflow-only change; there is no unit test for the release matrix. Verified by parsing both workflows as YAML and evaluating the expression for all six matrix entries (table above). ## Reviewer Notes Two questions worth a maintainer's call before this comes out of draft: 1. @jethachan — does SM121 need a warm AOT cache on **CUDA 12.9 aarch64** specifically, or is cu130 aarch64 sufficient? flashinfer-ai#3684's validation notes cover GB10 and RTX PRO 6000 but do not say which CUDA minor was used. 2. The alternative is to keep `12.1a` and instead land per-module filtering in the style of flashinfer-ai#3947. That needs a minor-version filter added to `get_nvcc_flags_list`, which today filters on major version only — more invasive than seems wise for v0.6.18.  ## Summary by CodeRabbit * **Bug Fixes** * Updated CUDA 12.9 wheel builds to use a consistent architecture set across platforms. * Preserved the specialized architecture target for CUDA 13.0 ARM64 builds. * Documented wheel-size constraints for affected builds.

yichengj0 requested review from aleozlx, bkryu, cyx-6, dhiraj113, jiahanc, jimmyzho, kahyunnam, nv-yunzheq, qiching, qsang-nv, saltyminty, samuellees, sricketts, yongwww, yyihuang and yzh119 as code owners July 13, 2026 17:59

flashinfer-bot added op: gemm op: moe labels Jul 13, 2026

yichengj0 mentioned this pull request Jul 13, 2026

ci: add 12.1a to aarch64 jit-cache wheel arch targets #3907

Closed

gemini-code-assist Bot reviewed Jul 13, 2026

View reviewed changes

yichengj0 force-pushed the dead-sm12x-cubins branch from 0f3c401 to 190aee6 Compare July 14, 2026 23:47

aleozlx reviewed Jul 15, 2026

View reviewed changes

aleozlx approved these changes Jul 15, 2026

View reviewed changes

aleozlx added the run-ci label Jul 15, 2026

yichengj0 force-pushed the dead-sm12x-cubins branch from 190aee6 to 4053636 Compare July 15, 2026 05:43

aleozlx approved these changes Jul 17, 2026

View reviewed changes

kahyunnam approved these changes Jul 17, 2026

View reviewed changes

bkryu merged commit 91054b9 into flashinfer-ai:main Jul 17, 2026
49 of 56 checks passed

yichengj0 mentioned this pull request Jul 20, 2026

perf(jit): prune arch gencode that dispatch can never load #4073

Merged

jimmyzho mentioned this pull request Jul 30, 2026

v0.6.16 release highlights draft #4278

Closed

kahyunnam mentioned this pull request Aug 14, 2026

WIP jit-cache wheel size fix: drop 12.1a from the cu129 aarch64 jit-cache wheel #4527

Merged

5 tasks

Uh oh!

Conversation

yichengj0 commented Jul 13, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

aleozlx Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

aleozlx Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

yichengj0 Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

aleozlx Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

yichengj0 Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

aleozlx commented Jul 15, 2026

Uh oh!

flashinfer-bot commented Jul 15, 2026

Uh oh!

flashinfer-bot commented Jul 15, 2026

Uh oh!

aleozlx commented Jul 17, 2026

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

bkryu commented Jul 17, 2026

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

bkryu commented Jul 17, 2026

Uh oh!

bkryu commented Jul 17, 2026

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

kahyunnam left a comment

Choose a reason for hiding this comment

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yichengj0 commented Jul 13, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 13, 2026 •

edited

Loading