Fix error when disabling new DeepGEMM #7198
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @fzyzcjy, 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 fixes an issue where certain DeepGEMM-specific configurations or operations were being applied unconditionally, leading to errors when the feature was explicitly disabled. The changes introduce checks to ensure these actions only occur when JIT DeepGEMM is enabled.
Highlights
- Bug Fix: I've addressed an error that occurred when the new DeepGEMM feature was disabled.
- Conditional Logic: Specific DeepGEMM-related operations and parameters are now guarded by a check to ensure JIT DeepGEMM is enabled before being applied.
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and 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 to provide feedback.
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
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configureGemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request correctly addresses an issue where features related to a new DeepGEMM version/type (gated by DEEPGEMM_V202506 and DEEPGEMM_SCALE_UE8M0) were active even when the main JIT DeepGEMM feature (ENABLE_JIT_DEEPGEMM) was disabled. The changes ensure these features are now conditional on ENABLE_JIT_DEEPGEMM being true, which should prevent the error mentioned in the PR title.
To further improve the PR:
- Please consider filling out the Motivation and Modifications sections in the PR description to provide more context about the error and the fix.
- It would be beneficial to add unit tests that specifically cover the scenario where
ENABLE_JIT_DEEPGEMMis false. These tests should verify that the error no longer occurs and that the system behaves as expected.
One minor code suggestion is made for python/sglang/srt/layers/moe/ep_moe/token_dispatcher.py to enhance semantic clarity regarding flag usage, though the current logic is functionally correct given the existing aliases.
| round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | ||
| and deep_gemm_wrapper.DEEPGEMM_V202506, | ||
| use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | ||
| and deep_gemm_wrapper.DEEPGEMM_V202506, |
There was a problem hiding this comment.
The logic to make round_scale and use_ue8m0 conditional on ENABLE_JIT_DEEPGEMM is correct and addresses the issue described in the PR title.
For improved semantic clarity and future maintainability, consider using deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 for the use_ue8m0 parameter. While DEEPGEMM_V202506 and DEEPGEMM_SCALE_UE8M0 are currently aliases (both True), they are used for distinct features elsewhere in the codebase (e.g., in python/sglang/srt/layers/moe/ep_moe/layer.py, DEEPGEMM_SCALE_UE8M0 is specifically related to scale_ue8m0). Aligning the flag usage here would make the code's intent clearer if these flags were to diverge in the future.
| round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | |
| and deep_gemm_wrapper.DEEPGEMM_V202506, | |
| use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | |
| and deep_gemm_wrapper.DEEPGEMM_V202506, | |
| round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | |
| and deep_gemm_wrapper.DEEPGEMM_V202506, | |
| use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | |
| and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, # Align with specific flag for ue8m0 features |
| if ( | ||
| deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM | ||
| and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 | ||
| ): |
There was a problem hiding this comment.
The condition for calling _weight_requant_ue8m0 now correctly includes deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM. This ensures that weight requantization for ue8m0 only occurs if JIT DeepGEMM is enabled, which aligns with the PR's goal to fix errors when DeepGEMM is disabled.
The use of parentheses (A and B) for the if condition is good for readability.
Motivation
local test on 8xB200 passes gsm8k
Modifications
Checklist