Revert "[1/2] sgl-kernel: Fuse routed scaling factor into select_experts" #8706
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @hnyls2002, 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 reverts a previous change that introduced a boolean parameter apply_routed_scaling_factor_on_output and its associated logic within the moe_fused_gate kernel. The original change aimed to fuse a routed scaling factor into expert selection, but this revert removes that specific functionality, restoring the code to its prior state.
Highlights
- Removed the apply_routed_scaling_factor_on_output boolean parameter: This boolean parameter was removed from the moe_fused_gate function signature across its C++ and Python definitions, including common_extension.cc, moe_fused_gate.cu, sgl_kernel_ops.h, and moe.py.
- Eliminated conditional scaling logic: The conditional logic within the CUDA kernel (moe_fused_gate_impl in moe_fused_gate.cu) that applied the routed_scaling_factor based on the now-removed boolean parameter has been eliminated.
- Updated function calls: All internal calls to the moe_fused_gate kernel functions (moe_fused_gate_impl, moe_fused_gate_kernel, moe_fused_gate_kernel_dynamic) were updated to no longer pass the removed apply_routed_scaling_factor_on_output parameter.
- Modified Python wrapper and test cases: The moe_fused_gate Python wrapper and its corresponding test cases in test_moe_fused_gate.py were modified to align with the updated function signature, removing the parameter from parametrization and function calls.
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 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 or fill out our survey 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 configure Gemini 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 reverts a previous change that fused a routed scaling factor into select_experts. The changes correctly remove the apply_routed_scaling_factor_on_output parameter and its associated logic from C++ kernels, Python bindings, and tests. My review focuses on ensuring the revert is clean and correct. I have one suggestion to improve the clarity of a docstring in the Python code that was updated as part of this revert.
| # num_fused_shared_experts: if > 0, the last several experts will be replaced with shared experts | ||
| # routed_scaling_factor: if > 0, the shared experts will be scaled by this factor |
There was a problem hiding this comment.
The docstring for routed_scaling_factor is misleading. It states that shared experts are 'scaled by this factor', but the implementation actually divides their weights by this factor. This could cause confusion. Let's update the docstrings for both parameters to be more precise.
| # num_fused_shared_experts: if > 0, the last several experts will be replaced with shared experts | |
| # routed_scaling_factor: if > 0, the shared experts will be scaled by this factor | |
| # num_fused_shared_experts: if > 0, the last several experts will be replaced with shared experts. | |
| # routed_scaling_factor: if > 0, the weights of shared experts will be divided by this factor. |
Reverts #8364