[sgl-kernel][5/N]Support Expert Specialization Grouped GEMM #12666
Conversation
HydraQYH
requested review from
BBuf,
Edwardf0t1,
HaiShaw,
Ying1123,
ch-wan,
ispobock,
kushanam,
merrymercy and
zhyncs
as code owners
Summary of ChangesHello @HydraQYH, 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 significantly enhances the Mixture-of-Experts (MoE) kernel performance by integrating support for Expert Specialization (ES) in CUTLASS grouped GEMM operations. It introduces a specialized kernel that leverages the capabilities of NVIDIA Hopper (SM90) GPUs, providing a performance boost for MoE computations. The changes are controlled by a new flag, allowing for flexible activation of these hardware-specific optimizations within the existing FP8 fused expert function. 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 introduces support for Expert Specialization Grouped GEMM in the CUTLASS MoE kernels, controlled by a new enable_es flag. The implementation is straightforward, but I've identified a few areas for improvement. There's significant code duplication in cutlass_moe.py that could be refactored to improve maintainability. More importantly, the new specialized kernel path is not being verified in the correctness tests, which is a high-priority issue to address. I've also noted that the test logic for enabling this feature is brittle due to hardcoded device names. Addressing these points will improve the quality and robustness of this new feature.
| if is_sm90_supported() and es_up: | ||
| es_fp8_blockwise_scaled_grouped_mm( | ||
| c1, | ||
| rep_a_q, | ||
| w1_q, | ||
| rep_a1_scales, | ||
| w1_scale, | ||
| a1_strides, | ||
| a1_strides, | ||
| c1_strides, | ||
| problem_sizes1, | ||
| expert_offsets[:-1], | ||
| workspace, | ||
| ) | ||
| else: | ||
| fp8_blockwise_scaled_grouped_mm( | ||
| c1, | ||
| a_ptrs, | ||
| b_ptrs, | ||
| out_ptrs, | ||
| a_scales_ptrs, | ||
| b_scales_ptrs, | ||
| rep_a_q, | ||
| w1_q, | ||
| rep_a1_scales, | ||
| w1_scale, | ||
| a1_strides, | ||
| a1_strides, | ||
| c1_strides, | ||
| a_sf_layout, | ||
| w_sf_layout, | ||
| problem_sizes1, | ||
| expert_offsets[:-1], | ||
| workspace, | ||
| ) | ||
|
|
||
| intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype) | ||
| silu_and_mul(c1, intermediate) | ||
|
|
||
| intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128) | ||
|
|
||
| fp8_blockwise_scaled_grouped_mm( | ||
| c2, | ||
| a_ptrs, | ||
| b_ptrs, | ||
| out_ptrs, | ||
| a_scales_ptrs, | ||
| b_scales_ptrs, | ||
| intemediate_q, | ||
| w2_q, | ||
| a2_scale, | ||
| w2_scale, | ||
| a2_strides, | ||
| a2_strides, | ||
| c2_strides, | ||
| a_sf_layout, | ||
| w_sf_layout, | ||
| problem_sizes2, | ||
| expert_offsets[:-1], | ||
| workspace, | ||
| ) | ||
| if is_sm90_supported() and es_down: | ||
| es_fp8_blockwise_scaled_grouped_mm( | ||
| c2, | ||
| intemediate_q, | ||
| w2_q, | ||
| a2_scale, | ||
| w2_scale, | ||
| a2_strides, | ||
| a2_strides, | ||
| c2_strides, | ||
| problem_sizes2, | ||
| expert_offsets[:-1], | ||
| workspace, | ||
| ) | ||
| else: | ||
| fp8_blockwise_scaled_grouped_mm( | ||
| c2, | ||
| a_ptrs, | ||
| b_ptrs, | ||
| out_ptrs, | ||
| a_scales_ptrs, | ||
| b_scales_ptrs, | ||
| intemediate_q, | ||
| w2_q, | ||
| a2_scale, | ||
| w2_scale, | ||
| a2_strides, | ||
| a2_strides, | ||
| c2_strides, | ||
| a_sf_layout, | ||
| w_sf_layout, | ||
| problem_sizes2, | ||
| expert_offsets[:-1], | ||
| workspace, | ||
| ) |
There was a problem hiding this comment.
There is significant code duplication between the GEMM calls for the up-projection (lines 162-196) and the down-projection (lines 203-237). Both blocks contain similar if/else logic to switch between the standard and expert specialization kernels. This redundancy makes the code harder to read and maintain.
To improve this, you could refactor the logic into a helper function. This function would encapsulate the conditional dispatch to either es_fp8_blockwise_scaled_grouped_mm or fp8_blockwise_scaled_grouped_mm.
Here is an example of how you could structure such a helper:
def _dispatch_gemm(is_es_enabled, out, a, b, a_scales, b_scales, a_strides, c_strides, problem_sizes, expert_offsets, workspace, **kwargs):
if is_sm90_supported() and is_es_enabled:
es_fp8_blockwise_scaled_grouped_mm(
out,
a,
b,
a_scales,
b_scales,
a_strides,
a_strides, # b_strides for GEMM
c_strides,
problem_sizes,
expert_offsets,
workspace,
)
else:
fp8_blockwise_scaled_grouped_mm(
out,
kwargs["a_ptrs"],
kwargs["b_ptrs"],
kwargs["out_ptrs"],
kwargs["a_scales_ptrs"],
kwargs["b_scales_ptrs"],
a,
b,
a_scales,
b_scales,
a_strides,
a_strides, # b_strides for GEMM
c_strides,
kwargs["a_sf_layout"],
kwargs["w_sf_layout"],
problem_sizes,
expert_offsets,
workspace,
)This would make the main function body much cleaner and avoid repeating the same logic.
2 participants
Motivation
Add the flag parameter to CUTLASS MoE to enable the expert specialization kernel.
Modifications
Accuracy Tests
check.log
Benchmarking and Profiling
Before using es:
After using es:
upstream_cutlass.log
upstream_es.log
Checklist