[apply][2/2] Fused qk_norm_rope for Qwen3-MoE

[yuan-luo]

Motivation

Qwen3-MoE qk_norm and rope ops account for a non-trivial proportion in model inference. See below:

Take Qwen3-30B-A3B for example, it has 48 Decode Layers. Each layer runs qk_norm and rope embedding ops. So this kernel enhancement can be amplified in linear along with the scale of the model layers.

This PR is to fuse qk_norm and rope into one kernel so as to improve performance.

Note:

• - sgl-kernel PR: [sgl-kernel][1/2] Fused qk_norm_rope for Qwen3-MoE #14036
• - apply srt PR: this PR.

Modifications

1. Add CUDA kernel for the fusion.
2. Revise Qwen3-MoE model to adapt for the new kernel.

Accuracy Tests

Without PR:
server:
python3 -m sglang.launch_server --model Qwen/Qwen3-30B-A3B --tp-size 4 --port 30000 --attention-backend flashinfer

gsm8k client result:
➜  sglang_dev git:(fused_qk_norm_rope) python3 benchmark/gsm8k/bench_sglang.py --num-questions 200 --parallel 128 --num-shots 8 --port 30010
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:10<00:00, 18.62it/s]
Accuracy: 0.730
Invalid: 0.000
Latency: 10.876 s
Output throughput: 5015.168 token/s

With PR:
server:
python3 -m sglang.launch_server --model Qwen/Qwen3-30B-A3B --tp-size 4 --port 30000 --attention-backend flashinfer --enable-fused-qk-norm-rope

gsm8k client:
➜  sglang_dev git:(fused_qk_norm_rope) python3 benchmark/gsm8k/bench_sglang.py --num-questions 200 --parallel 128 --num-shots 8 --port 30010
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:09<00:00, 21.31it/s]
Accuracy: 0.740
Invalid: 0.000
Latency: 9.483 s
Output throughput: 5529.482 token/s

Benchmarking and Profiling

Test Qwen3-235B-A22B on 8xH20.
TTFT: 6.0% drop from 990.23ms to 935.69ms

Server:

python3 -m sglang.launch_server --model /home/admin/Qwen3-235B-A22B \
  --tp-size 8 --port 30000 --attention-backend flashinfer \
  --disable-radix-cache \
  # --enable-fused-qk-norm-rope

Benchmark client:

python3 -m sglang.bench_serving --backend sglang \
  --num-prompts 500 --dataset-name random \
  --random-input 500 --random-output 100 \
  --port 30000 --random-range-ratio 1.0 \
  --max-concurrency 32 \
  --dataset-path /tmp/ShareGPT_V3_unfiltered_cleaned_split.json

PR:
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 32        
Successful requests:                     500       
Benchmark duration (s):                  55.29     
Total input tokens:                      250000    
Total input text tokens:                 250000    
Total input vision tokens:               0         
Total generated tokens:                  50000     
Total generated tokens (retokenized):    50000     
Request throughput (req/s):              9.04      
Input token throughput (tok/s):          4521.59   
Output token throughput (tok/s):         904.32    
Peak output token throughput (tok/s):    1344.00   
Peak concurrent requests:                64        
Total token throughput (tok/s):          5425.91   
Concurrency:                             31.34     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   3465.24   
Median E2E Latency (ms):                 3470.12   
---------------Time to First Token----------------
Mean TTFT (ms):                          935.69    
Median TTFT (ms):                        903.84    
P99 TTFT (ms):                           1580.97   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          25.55     
Median TPOT (ms):                        25.73     
P99 TPOT (ms):                           31.25     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           25.55     
Median ITL (ms):                         24.20     
P95 ITL (ms):                            25.02     
P99 ITL (ms):                            28.21     
Max ITL (ms):                            1129.26   
==================================================

Main:
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 32        
Successful requests:                     500       
Benchmark duration (s):                  55.69     
Total input tokens:                      250000    
Total input text tokens:                 250000    
Total input vision tokens:               0         
Total generated tokens:                  50000     
Total generated tokens (retokenized):    50000     
Request throughput (req/s):              8.98      
Input token throughput (tok/s):          4488.99   
Output token throughput (tok/s):         897.80    
Peak output token throughput (tok/s):    1376.00   
Peak concurrent requests:                64        
Total token throughput (tok/s):          5386.78   
Concurrency:                             31.35     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   3491.58   
Median E2E Latency (ms):                 3452.26   
---------------Time to First Token----------------
Mean TTFT (ms):                          990.23    
Median TTFT (ms):                        994.89    
P99 TTFT (ms):                           1672.83   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          25.27     
Median TPOT (ms):                        25.39     
P99 TPOT (ms):                           31.14     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           25.27     
Median ITL (ms):                         23.83     
P95 ITL (ms):                            24.61     
P99 ITL (ms):                            27.60     
Max ITL (ms):                            1163.50   
==================================================

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

Hello @yuan-luo, I'm Gemini Code Assist¹! 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 delivers a crucial performance enhancement for the Qwen3-MoE model by integrating a fused CUDA kernel for query-key normalization and Rotary Position Embedding. This optimization directly addresses a significant computational overhead, resulting in faster inference and higher throughput. The changes involve low-level kernel implementation, model adaptation, and a new configuration option, all validated by comprehensive unit tests.

Highlights

• Performance Optimization: This pull request introduces a significant performance optimization for Qwen3-MoE models by fusing the qk_norm (query-key normalization) and rope (Rotary Position Embedding) operations into a single, highly optimized CUDA kernel. This fusion targets a bottleneck identified in Qwen3-MoE inference, where these operations account for a large proportion of computational time.
• Improved Inference Speed: Benchmarking results demonstrate a notable improvement in inference speed. On the gsm8k client, latency decreased from 13.911 seconds to 11.701 seconds, and output throughput increased from 4008.419 tokens/s to 4489.955 tokens/s when the fused kernel was enabled.
• New CUDA Kernel and Integration: A new CUDA kernel, fused_qknorm_rope_kernel.cu, has been added to perform the combined normalization and rotary embedding. The Qwen3-MoE model's attention mechanism has been revised to utilize this new kernel, with a fallback to the non-fused implementation if the fused kernel is not enabled or compatible.
• Configurable Feature: A new command-line argument, --enable-fused-qk-norm-rope, has been added to the server arguments, allowing users to explicitly enable or disable this fused operation. This provides flexibility for deployment and testing.
• Comprehensive Testing: A new test file, test_fused_qk_norm_rope.py, has been added to thoroughly verify the correctness of the fused kernel. It compares the output of the fused operation against a PyTorch reference implementation across various head dimensions, head group configurations, token counts, and RoPE styles (NeoX and interleaved).

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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. 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

1. 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. ↩

[yuan-luo]

The result is as expected.

➜  bench_script python test_openai.py
ChatCompletion(id='ea37920145974ee9888dead212b67864', choices=[Choice(finish_reason='length', index=0, logprobs=None, message=ChatCompletionMessage(content="<think>\nOkay, the user asked for three countries and their capitals, and then how I rank them. Let me start by picking three countries. Maybe the US, Japan, and Brazil. Their capitals are Washington, D.C., Tokyo, and Brasília. Now, how to rank them? The user didn't specify the criteria, so I need to think of possible ways. Maybe by population, economic size, or cultural influence. Let me check the population. The US has around 330 million, Japan about 125 million, Brazil 215 million. So US first, Brazil second, Japan third. But if I consider GDP, the US is the largest, then Japan, then Brazil. Alternatively, cultural influence: Japan has a strong cultural impact, maybe higher than Brazil. But the user might not have a specific criteria. I should mention that the ranking depends on the criteria and provide examples. Also, make sure the capitals are correct. Washington, D.C", refusal=None, role='assistant', annotations=None, audio=None, function_call=None, tool_calls=None, reasoning_content=None), matched_stop=None)], created=1764167050, model='default', object='chat.completion', service_tier=None, system_fingerprint=None, usage=CompletionUsage(completion_tokens=200, prompt_tokens=33, total_tokens=233, completion_tokens_details=None, prompt_tokens_details=None, reasoning_tokens=0), metadata={'weight_version': 'default'})

[gemini-code-assist]

Code Review

This pull request introduces a fused kernel for QK normalization and RoPE, which is a great performance optimization for Qwen3-MoE models. The implementation is well-structured and adapted from TensorRT-LLM. I have a few suggestions for improving maintainability, such as removing a redundant import and some commented-out code. I also noticed that the new kernel is placed in MoE-related directories, which might not be the most intuitive location. More importantly, I've found a critical issue in the new test file where the reference implementation is incorrect, and a gap in test coverage for the YaRN scaling feature. Addressing these points will ensure the new kernel is both correct and robust.

[ispobock]

Could you split the kernel and python parts into two PRs?

[yuan-luo]

Could you split the kernel and python parts into two PRs?

Sure, will do it tmr.

[Kangyan-Zhou]

/tag-and-rerun-ci

[yuan-luo]

Could you split the kernel and python parts into two PRs?

Done. #14036

[BBuf]

Line 186-199 and 216-229 are duplicate code?

[yuan-luo]

Refactored.

[BBuf]

Please delete model modify in sgl-kernel pr.

[yuan-luo]

Please delete model modify in sgl-kernel pr.

This PR is the model modify. sgl-kernel is in #14036. I modified cuda code on both PRs.

[yuan-luo]

Please delete model modify in sgl-kernel pr.
done.

[yuan-luo]

We need to wait for the new sgl-kernel release before merging this PR.

[HydraQYH]

LGTM

[yuan-luo]

Tested on H20 for Qwen3-235B-A22B, obtained 6% TTFT improvement.
TTFT: drop from 990.23ms to 935.69ms

python3 -m sglang.launch_server --model /home/admin/Qwen3-235B-A22B \
  --tp-size 8 --port 30000 --attention-backend flashinfer \
  --disable-radix-cache \
  --enable-fused-qk-norm-rope

python3 -m sglang.bench_serving --backend sglang \
  --num-prompts 500 --dataset-name random \
  --random-input 500 --random-output 100 \
  --port 30000 --random-range-ratio 1.0 \
  --max-concurrency 32 \
  --dataset-path /tmp/ShareGPT_V3_unfiltered_cleaned_split.json

PR:
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 32        
Successful requests:                     500       
Benchmark duration (s):                  55.29     
Total input tokens:                      250000    
Total input text tokens:                 250000    
Total input vision tokens:               0         
Total generated tokens:                  50000     
Total generated tokens (retokenized):    50000     
Request throughput (req/s):              9.04      
Input token throughput (tok/s):          4521.59   
Output token throughput (tok/s):         904.32    
Peak output token throughput (tok/s):    1344.00   
Peak concurrent requests:                64        
Total token throughput (tok/s):          5425.91   
Concurrency:                             31.34     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   3465.24   
Median E2E Latency (ms):                 3470.12   
---------------Time to First Token----------------
Mean TTFT (ms):                          935.69    
Median TTFT (ms):                        903.84    
P99 TTFT (ms):                           1580.97   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          25.55     
Median TPOT (ms):                        25.73     
P99 TPOT (ms):                           31.25     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           25.55     
Median ITL (ms):                         24.20     
P95 ITL (ms):                            25.02     
P99 ITL (ms):                            28.21     
Max ITL (ms):                            1129.26   
==================================================

Main:
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 32        
Successful requests:                     500       
Benchmark duration (s):                  55.69     
Total input tokens:                      250000    
Total input text tokens:                 250000    
Total input vision tokens:               0         
Total generated tokens:                  50000     
Total generated tokens (retokenized):    50000     
Request throughput (req/s):              8.98      
Input token throughput (tok/s):          4488.99   
Output token throughput (tok/s):         897.80    
Peak output token throughput (tok/s):    1376.00   
Peak concurrent requests:                64        
Total token throughput (tok/s):          5386.78   
Concurrency:                             31.35     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   3491.58   
Median E2E Latency (ms):                 3452.26   
---------------Time to First Token----------------
Mean TTFT (ms):                          990.23    
Median TTFT (ms):                        994.89    
P99 TTFT (ms):                           1672.83   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          25.27     
Median TPOT (ms):                        25.39     
P99 TPOT (ms):                           31.14     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           25.27     
Median ITL (ms):                         23.83     
P95 ITL (ms):                            24.61     
P99 ITL (ms):                            27.60     
Max ITL (ms):                            1163.50   
==================================================