[Optimization] Accelerate Qwen3 QK RMSNorm via Fused Triton Kernel

[Sunny-bot1]

Motivation

Qwen3 系列模型在执行 QK RMSNorm 前，需要先对 qkv_out 进行 split，归一化后再 concat 回去。这一过程涉及多次张量重排与内存拷贝，开销较大，同时现有的 RMSNorm kernel 在性能上也并非最优。

Modifications

通过 Triton kernel 融合优化，利用 offset 直接索引 Q/K/V 数据，把 split → reshape → RMSNorm(Q) → reshape → RMSNorm(K) → concat 融合成一个 Triton kernel，消除了 split/concat 开销，并提升了 RMSNorm kernel 的性能。

kernel性能

Qwen3-30B-A3B TP1

======== qk_norm_triton_fused Performance ========
Batch Size      Last 5 Times (us)                        Last Time (us)  TFlops          TB/s           
1               [1.88 1.88 1.88 1.88 1.88]               1.88           
8               [2.08 2.08 2.08 2.08 2.08]               2.08           
64              [2.29 2.29 2.29 2.29 2.29]               2.29           
128             [2.71 2.71 2.71 2.71 2.71]               2.71           
1024            [7.08 7.08 7.08 7.08 7.08]               7.08           
2048            [12.08 12.08 12.08 12.08 12.08]          12.08          
4096            [23.54 23.33 23.33 23.54 23.54]          23.54          
8192            [52.92 52.92 52.92 52.92 52.92]          52.92       
.======== qkv_split + qk_norm_paddle + qkv_concat Performance ========
Batch Size      Last 5 Times (us)                        Last Time (us)  TFlops          TB/s           
1               [8.96 8.96 8.96 8.96 8.96]               8.96           
8               [8.96 8.96 8.96 8.96 8.96]               8.96           
64              [10.62 10.62 10.62 10.62 10.62]          10.62          
128             [12.92 12.92 12.92 12.92 12.92]          12.92          
1024            [55.42 55.42 55.42 55.21 55.42]          55.42          
2048            [108.12 108.12 108.33 108.12 108.12]     108.12         
4096            [203.96 203.96 203.96 203.96 204.17]     204.17         
8192            [393.54 393.75 393.96 393.75 393.75]     393.75

Qwen3-235B-A22B TP4

======== qk_norm_triton_fused Performance ========
Batch Size      Last 5 Times (us)                        Last Time (us)  TFlops          TB/s           
1               [1.91 1.91 1.91 1.91 1.91]               1.91           
8               [2.02 2.02 2.02 2.02 2.02]               2.02           
64              [2.13 2.13 2.13 2.13 2.13]               2.13           
128             [2.34 2.34 2.34 2.34 2.34]               2.34           
1024            [4.57 4.57 4.57 4.57 4.57]               4.57           
2048            [7.02 7.02 7.02 7.02 7.02]               7.02           
4096            [12.02 12.02 12.02 12.02 12.02]          12.02          
8192            [22.02 22.02 22.02 22.02 22.02]          22.02          
.======== qkv_split + qk_norm_paddle + qkv_concat Performance ========
Batch Size      Last 5 Times (us)                        Last Time (us)  TFlops          TB/s           
1               [8.72 8.72 8.72 8.72 8.72]               8.72           
8               [9.36 9.36 9.36 9.36 9.36]               9.36           
64              [10.21 10.21 10.21 10.21 10.21]          10.21          
128             [11.49 11.49 11.49 11.49 11.49]          11.49          
1024            [27.34 27.34 27.23 27.23 27.23]          27.23          
2048            [53.51 53.51 53.62 53.62 53.51]          53.51          
4096            [103.19 103.3  103.19 103.3  103.19]     103.19         
8192            [196.81 196.81 196.7  196.81 196.7 ]     196.7 

GLM_4.6 TP4

======== qk_norm_triton_fused Performance ========
Batch Size      Last 5 Times (us)                        Last Time (us)  TFlops          TB/s           
1               [1.85 1.85 1.85 1.85 1.85]               1.85           
8               [2.07 2.07 2.07 2.07 2.07]               2.07           
64              [2.17 2.17 2.17 2.17 2.17]               2.17           
128             [2.39 2.39 2.39 2.39 2.39]               2.39           
1024            [5.87 5.87 5.87 5.87 5.87]               5.87           
2048            [9.46 9.46 9.46 9.46 9.46]               9.46           
4096            [16.96 16.96 16.96 16.96 16.96]          16.96          
8192            [38.8 38.7 38.8 38.8 38.8]               38.8           
.======== qkv_split + qk_norm_paddle + qkv_concat Performance ========
Batch Size      Last 5 Times (us)                        Last Time (us)  TFlops          TB/s           
1               [8.8 8.8 8.8 8.8 8.8]                    8.8            
8               [8.8 8.8 8.8 8.8 8.8]                    8.8            
64              [10.43 10.43 10.43 10.43 10.43]          10.43          
128             [11.63 11.63 11.63 11.63 11.63]          11.63          
1024            [39.02 39.13 39.13 39.02 39.13]          39.13          
2048            [78.7 78.7 78.7 78.7 78.7]               78.7           
4096            [150.11 150.22 150.22 150.22 150.22]     150.22         
8192            [289.78 289.67 289.67 289.78 289.67]     289.67

模型性能

Qwen3-30B-A3B FP8下TPS提升5%🚀

	并发	输入长度	输出长度	解码速度	QPS	TPS
优化前	64	2025.80	830.86	41.23	3.062	8747.04
优化后	64	2025.80	823.50	42.36	3.232	9208.39

Usage or Command

Accuracy Tests

已添加单测，相比paddle api误差在1e-4内。

Checklist

• Add at least a tag in the PR title.
  • Tag list: [[FDConfig],[APIServer],[Engine], [Scheduler], [PD Disaggregation], [Executor], [Graph Optimization], [Speculative Decoding], [RL], [Models], [Quantization], [Loader], [OP], [KVCache], [DataProcessor], [BugFix], [Docs], [CI], [Optimization], [Feature], [Benchmark], [Others], [XPU], [HPU], [GCU], [DCU], [Iluvatar], [Metax]]
  • You can add new tags based on the PR content, but the semantics must be clear.
• Format your code, run pre-commit before commit.
• Add unit tests. Please write the reason in this PR if no unit tests.
• Provide accuracy results.
• If the current PR is submitting to the release branch, make sure the PR has been submitted to the develop branch, then cherry-pick it to the release branch with the [Cherry-Pick] PR tag.

[paddle-bot]

Thanks for your contribution!

[codecov-commenter]

Codecov Report

❌ Patch coverage is 54.45545% with 46 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@60ee72f). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...executor/ops/triton_ops/qk_rmsnorm_fused_kernel.py	38.09%	26 Missing ⚠️
fastdeploy/model_executor/layers/normalization.py	61.76%	11 Missing and 2 partials ⚠️
fastdeploy/model_executor/models/glm4_moe.py	16.66%	4 Missing and 1 partial ⚠️
fastdeploy/model_executor/models/qwen3.py	87.50%	1 Missing ⚠️
...stdeploy/model_executor/ops/triton_ops/__init__.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop    #5880   +/-   ##
==========================================
  Coverage           ?   67.33%           
==========================================
  Files              ?      348           
  Lines              ?    44688           
  Branches           ?     6875           
==========================================
  Hits               ?    30092           
  Misses             ?    12374           
  Partials           ?     2222           

Flag	Coverage Δ
GPU	67.33% <54.45%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.