fix: Add fused MOE and GEMM AOT modules for SM121#2654
fix: Add fused MOE and GEMM AOT modules for SM121#2654yzh119 merged 2 commits intoflashinfer-ai:mainfrom
Conversation
SM121 (DGX Spark / GB10) was missing fused_moe, gemm, and fp4_gemm_cutlass modules in the AOT build. Only fp4_quantization was generated for SM121 systems. The JIT runtime path already correctly dispatches SM121 to SM120 modules (both share the same CUTLASS kernels), but the AOT pre- compilation path under `if has_sm121:` only included gen_fp4_quantization_sm121_module(). Reuse the SM120 module generators for SM121 since they use supported_major_versions=[12] which compiles for all SM12x targets. The dedup at the end of gen_all_modules() prevents duplicates when both has_sm120 and has_sm121 are True. Validated on DGX Spark (SM121a, NVIDIA GB10): Before fix — SM121 AOT modules: fp4_quantization_121 Total SM12x modules: 1 After fix — SM121 AOT modules: fp4_quantization_121 fused_moe_120 gemm_sm120 fp4_gemm_cutlass_sm120 Total SM12x modules: 4 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
blake-snc
requested review from
bkryu,
jimmyzho,
kahyunnam,
nv-yunzheq and
yzh119
as code owners
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a gap in the Ahead-Of-Time (AOT) compilation process for SM121 systems. Previously, critical fused MOE and GEMM modules were missing from SM121 AOT builds, leading to runtime JIT compilation. By integrating the appropriate SM120 module generators, this change ensures that SM121 systems benefit from pre-compiled kernels, thereby improving performance and reducing startup overhead for applications utilizing these operations. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughgen_all_modules in Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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}
There was a problem hiding this comment.
Code Review
This pull request correctly adds the missing AOT (Ahead-of-Time) compilation modules for SM121 by reusing the existing SM120 module generators. The logic is sound, as SM121 shares the same CUTLASS kernels as SM120. My review includes one suggestion to refactor the code to eliminate duplication, which would improve the long-term maintainability of this file.
flashinfer/aot.py
Outdated
| # SM121 shares the same CUTLASS kernels as SM120 for fused MOE and GEMM. | ||
| # The SM120 module generators use supported_major_versions=[12] which | ||
| # compiles for all SM12x targets. Dedup at end of gen_all_modules() | ||
| # prevents duplicate modules when both has_sm120 and has_sm121 are True. | ||
| jit_specs.append(gen_cutlass_fused_moe_sm120_module()) | ||
| jit_specs.append(gen_gemm_sm120_module()) | ||
| jit_specs.append(gen_gemm_sm120_module_cutlass_fp4()) |
There was a problem hiding this comment.
This change correctly adds the necessary modules for SM121. However, it introduces code duplication, as lines 524-526 are identical to lines 515-517 within the if has_sm120: block.
To improve maintainability and avoid this duplication, consider refactoring this logic. Since sm120 and sm121 share these modules, you can group their addition under a combined condition.
Here's a suggested refactoring:
if has_sm120:
jit_specs.append(gen_fp4_quantization_sm120_module())
if has_sm121:
jit_specs.append(gen_fp4_quantization_sm121_module())
if has_sm120 or has_sm121:
# SM121 shares the same CUTLASS kernels as SM120 for fused MOE and GEMM.
# The SM120 module generators use supported_major_versions=[12] which
# compiles for all SM12x targets.
jit_specs.append(gen_cutlass_fused_moe_sm120_module())
jit_specs.append(gen_gemm_sm120_module())
jit_specs.append(gen_gemm_sm120_module_cutlass_fp4())This approach eliminates redundancy and makes the code's intent clearer. It also avoids relying on the dedup step at the end of the function for this particular case.
Per review feedback, refactor the SM120/SM121 fused MOE and GEMM module generation to use `if has_sm120 or has_sm121:` instead of duplicating the three calls in both blocks. The fp4_quantization modules remain separate since they use arch-specific nvcc flags. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
| # SM120 and SM121 share the same CUTLASS kernels for fused MOE and GEMM. | ||
| # The SM120 module generators use supported_major_versions=[12] which | ||
| # compiles for all SM12x targets. | ||
| jit_specs.append(gen_cutlass_fused_moe_sm120_module()) |
There was a problem hiding this comment.
probably better to rename them to sm_12x
## Summary
- Add `gen_cutlass_fused_moe_sm120_module()`, `gen_gemm_sm120_module()`,
and `gen_gemm_sm120_module_cutlass_fp4()` to the `if has_sm121:` block
in `aot.py`
- SM121 (DGX Spark / GB10) AOT builds now include fused MOE and GEMM
modules
## Problem
SM121 was missing fused_moe, gemm, and fp4_gemm_cutlass modules in the
AOT pre-compilation path. Only `gen_fp4_quantization_sm121_module()` was
generated for SM121 systems, while the JIT runtime path already
correctly dispatches SM121 to SM120 modules for all three operations.
This means AOT-prebuilt packages (e.g. `flashinfer-jit-cache`) targeting
SM121 would be missing fused MOE and GEMM kernels, forcing fallback to
JIT compilation at runtime.
## Fix
Reuse the SM120 module generators under `if has_sm121:` since SM120 and
SM121 share the same CUTLASS kernels. The SM120 generators already use
`supported_major_versions=[12]` which compiles for all SM12x targets.
The dedup at the end of `gen_all_modules()` prevents duplicate modules
when both `has_sm120` and `has_sm121` are True.
## Validation (DGX Spark, SM121a)
```
=== SM capabilities detected ===
sm121: True
=== BEFORE fix: SM121 AOT modules ===
fp4_quantization_121
Total SM12x modules: 1
=== AFTER fix: SM121 AOT modules ===
fp4_quantization_121
fused_moe_120
gemm_sm120
fp4_gemm_cutlass_sm120
Total SM12x modules: 4
```
JIT runtime dispatch confirmed: `get_cutlass_fused_moe_module("121")`,
`get_gemm_sm120_module()`, and `get_cutlass_fp4_gemm_module(12, 1)` all
correctly route SM121 to the SM120 modules at runtime.
Contributed by Second Nature Computing (https://joinsecondnature.com)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Chores**
* Improved SM120/SM121 support: generation of fused MOE, GEMM, and
FP4-quantization kernel modules now occurs under a unified condition for
SM120 or SM121, with FP4 quantization included alongside other SM121
modules earlier in the pipeline. Automatic deduplication remains to
prevent redundant kernels.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
2 participants
Summary
gen_cutlass_fused_moe_sm120_module(),gen_gemm_sm120_module(), andgen_gemm_sm120_module_cutlass_fp4()to theif has_sm121:block inaot.pyProblem
SM121 was missing fused_moe, gemm, and fp4_gemm_cutlass modules in the AOT pre-compilation path. Only
gen_fp4_quantization_sm121_module()was generated for SM121 systems, while the JIT runtime path already correctly dispatches SM121 to SM120 modules for all three operations.This means AOT-prebuilt packages (e.g.
flashinfer-jit-cache) targeting SM121 would be missing fused MOE and GEMM kernels, forcing fallback to JIT compilation at runtime.Fix
Reuse the SM120 module generators under
if has_sm121:since SM120 and SM121 share the same CUTLASS kernels. The SM120 generators already usesupported_major_versions=[12]which compiles for all SM12x targets. The dedup at the end ofgen_all_modules()prevents duplicate modules when bothhas_sm120andhas_sm121are True.Validation (DGX Spark, SM121a)
JIT runtime dispatch confirmed:
get_cutlass_fused_moe_module("121"),get_gemm_sm120_module(), andget_cutlass_fp4_gemm_module(12, 1)all correctly route SM121 to the SM120 modules at runtime.Contributed by Second Nature Computing (https://joinsecondnature.com)
Summary by CodeRabbit