Fix/moe_sm110 (to be tested) #2183
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughThis change modifies the MOE GEMM kernel dispatch logic for NVIDIA TensorRT-LLM, specifically narrowing the compute capability window for SM100 dispatch, expanding SM100+ fallback conditions to include SM110, and adding arch-specific gating for tile shape support checks to differentiate SM110 from other SM100 configurations. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🧠 Learnings (1)📚 Learning: 2025-11-12T03:35:17.583ZApplied to files:
🧬 Code graph analysis (1)csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h (2)
🔇 Additional comments (2)
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 |
Summary of ChangesHello @aleozlx, 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 functional regression in the Mixture-of-Experts (MoE) Cutlass backend for SM110 architectures, which previously failed to find viable configurations. The fix involves adjusting the kernel dispatch logic and relaxing certain cluster shape constraints specifically for SM110, ensuring it correctly utilizes the MoE kernels. These changes were necessary due to recent updates intended for newer architectures that inadvertently affected SM110. Highlights
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
|
![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 addresses a functional regression for SM110 in the MoE CUTLASS backend by adjusting the dispatch logic. The changes correctly move SM110 to a different dispatch path, aligning it with SM90 and SM120+ instead of SM100/103. A corresponding change is made to the tile shape validation to remove a failing constraint for SM110.
While the changes appear to correctly fix the issue, the patch to the validation logic reveals a structural inconsistency between how architectures are grouped for dispatch versus validation. I've left a comment with a suggestion to address this for better long-term maintainability.
| if constexpr (Arch::kMinComputeCapability != 110) { | ||
| // We use a runtime cluster shape for SM100, so we only support 1x1x1 and 2x1x1 cluster shapes. | ||
| if (cute::size<0>(ClusterShape{}) > 2 || cute::size<1>(ClusterShape{}) != 1 || | ||
| cute::size<2>(ClusterShape{}) != 1) { | ||
| return false; | ||
| } | ||
| } |
There was a problem hiding this comment.
While this change fixes the regression for SM110, special-casing it within are_tile_shapes_supported_sm100 highlights a structural inconsistency that could affect maintainability.
The changes in dispatchMoeGemmFinalDispatchTmaWarpSpecialized correctly group SM110 with SM90 and SM120+, separating it from the SM100/103 path. However, the validation logic in the calling function are_tile_shapes_supported still groups SM110 with SM100/103, which necessitates this patch.
This discrepancy makes the code harder to reason about, as the logic for SM110 is fragmented. For better long-term code health, the validation paths in are_tile_shapes_supported should be refactored to align with the dispatch paths. I recommend creating a follow-up technical debt issue to address this.
There was a problem hiding this comment.
agreed with Gemini's suggestion
yzh119
requested review from
djmmoss,
wenscarl,
yongwww and
yzh119
as code owners
|
/bot run |
|
[FAILED] Pipeline #40898746: 1/20 passed |
3 participants
📌 Description
Fix sm110 moe (cutlass backend) functional regression : no viable configs
Found dispatch logic changed and extra cluster constraint from #2020 / #1925 trying to support sm103 sm120 new features. The logic haven't been tested on sm110 due to a current lack of CI resources
(to be tested)
🔍 Related Issues
🚀 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.).Reviewer Notes
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.