optimize gdn gating and fused_qkvzba_split_reshape_cat#306
optimize gdn gating and fused_qkvzba_split_reshape_cat#306iforgetmyname merged 1 commit intosgl-project:mainfrom
Conversation
Summary of ChangesHello @RuixuanZhang06, 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 focuses on enhancing the performance of specific computational blocks within the 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. 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 optimizations for GDN gating and a utility function for tensor manipulation. My review focuses on improving performance, correctness, and code clarity.
In fused_gdn_gating.py, I've identified a potential performance bottleneck in the Triton kernel's grid configuration that could lead to serial processing of batches. I've suggested a more robust grid setup and corresponding kernel changes to ensure proper parallelization. I also pointed out minor issues like an import inside a function and confusing output tensor shapes.
In utils.py, the new fused_qkvzba_split_reshape_cat_torch function is well-structured. I've suggested using torch.split to make the tensor splitting logic more concise and readable.
Overall, the changes are good, but addressing the performance issue in the Triton kernel is important.
RuixuanZhang06
force-pushed
the
qwen3-next
branch
from
c64089e to
dbe3db6
Compare
RuixuanZhang06
force-pushed
the
qwen3-next
branch
from
dbe3db6 to
1a9c761
Compare
* upstream/main: optimize gdn gating and fused_qkvzba_split_reshape_cat (sgl-project#306) fix layout numTokensPerExpertTensor partial Initialization bug (sgl-project#303) Supplement A2 doc, software and hardware compatibility info (sgl-project#294) Added an environment variable to control whether to enable the Combine Ant Migration feature. (sgl-project#304) Support build with cann 8.5 (sgl-project#283) LoRA: Optimization LoRA kernels and refactoring (sgl-project#284) fix a2 single combine aclnn params Resolving the UB out-of-bounds issue caused by A2 dual-machine mixed operation (sgl-project#288) fix notify magic auto-increment bug (sgl-project#291) split_qkv_rmsnorm_rope bugfix (sgl-project#290) Optimize prepare_lens by removing device transfer (sgl-project#289)
…pu-old into bugfix * 'a3_topk-1' of https://github.com/luanyundu/sgl-kernel-npu-old: fix dispatch_layout to support topk -1 feature optimize gdn gating and fused_qkvzba_split_reshape_cat (sgl-project#306) fix layout numTokensPerExpertTensor partial Initialization bug (sgl-project#303) Supplement A2 doc, software and hardware compatibility info (sgl-project#294) Added an environment variable to control whether to enable the Combine Ant Migration feature. (sgl-project#304) Support build with cann 8.5 (sgl-project#283) LoRA: Optimization LoRA kernels and refactoring (sgl-project#284) fix a2 single combine aclnn params Resolving the UB out-of-bounds issue caused by A2 dual-machine mixed operation (sgl-project#288) fix notify magic auto-increment bug (sgl-project#291) split_qkv_rmsnorm_rope bugfix (sgl-project#290) Optimize prepare_lens by removing device transfer (sgl-project#289) Fix the performance degradation issue of the single-wheel operation in Ant Moving. (sgl-project#287) modify split_qkv_rmsnorm_rope (sgl-project#282)
* upstream/main: fix little batchsize and int8 quant on ci (sgl-project#302) optimize sinks attention (sgl-project#260) add swiglu_oai_triton (sgl-project#270) update tag to 2026.01.12 (sgl-project#312) feat:add performance compare (sgl-project#311) support add_gemma_rms_norm (sgl-project#310) optimize gdn gating and fused_qkvzba_split_reshape_cat (sgl-project#306) fix layout numTokensPerExpertTensor partial Initialization bug (sgl-project#303) Supplement A2 doc, software and hardware compatibility info (sgl-project#294) Added an environment variable to control whether to enable the Combine Ant Migration feature. (sgl-project#304)
2 participants
No description provided.