[ROCm][Perf] Fuse SWA q/kv RMSNorm and q FP8 group quant for DeepSeek-V4 by shen-shanshan · Pull Request #53540 · vllm-project/vllm

shen-shanshan · 2026-08-24T06:46:48Z

Purpose

On the ROCm path of DeepSeek-V4, every decode step's SWA token-insertion pipeline runs:

vLLM's Triton fused_q_kv_rmsnorm (norms the q-lora and kv latents to bf16).
aiter per-1x128 dynamic quant of the bf16 q-lora (input quant of the block-scaled FP8 wq_b GEMM).
the same quant a second time in the indexer's wq_b, which shares the q-lora with the attention layer.

This PR replaces 1+2 with a single aiter HIP kernel (fused_qk_rmsnorm_group_quant): it RMSNorms both latents (fp32 accumulate) and quantizes the q latent (per-1x128, fp32 scales) in one pass, so:

the hot path drops from 2 kernel launches to 1.
the bf16 q-lora write+read round-trip disappears; quantization now reads the fp32 accumulator directly (single rounding, equivalent-or-better precision).
the attention wq_b and the indexer wq_b both consume the pre-quantized (qr_fp8, qr_scale) pair and skip their own input quant (removes a third quant launch on the long-context path).
kv stays bf16 and feeds the existing fused insert kernel (fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert) unchanged.

Note

The change is scoped to the ROCm/AITER path and does not change the CUDA or XPU implementations.

Before fusing (4µs 839ns + 4µs 719ns + 4µs 719ns = 14 µs 277 ns):

After fusing (4µs 919ns, 2.90x faster):

Test Plan

Benchmark with SA InferenceX 8k1k workload.
Accuracy test with gsm8k dataset.

Test Result

Note

The tests below are executed on AMD MI350X GPU, so the performance could be slightly lower than the MI355X baseline, so I verified this PR with an A/B test compared with vLLM main on the same machine.

Benchmark

Concurrency	Version	Output token throughput (tok/s)	Change	Mean TTFT (ms)	Change	Mean TPOT (ms)	Change
1	Main	56.29	—	418.93	—	17.33	—
1	This PR	56.70	0.73% ↑	462.05	10.29% ↑	17.16	0.98% ↓
4	Main	181.61	—	514.31	—	20.95	—
4	This PR	182.49	0.48% ↑	517.49	0.62% ↑	20.84	0.53% ↓
16	Main	530.80	—	806.76	—	28.30	—
16	This PR	535.25	0.84% ↑	808.93	0.27% ↑	28.05	0.88% ↓
64	Main	1062.02	—	1780.55	—	57.27	—
64	This PR	1067.39	0.51% ↑	1778.38	0.12% ↓	56.98	0.51% ↓

Accuracy Test

Full GSM8K evaluation on this change (1,319 examples):

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9431	±	0.0064
		strict-match	5	exact_match	↑	0.9431	±	0.0064

Both results exceed the required 94% threshold.

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

tjtanaa · 2026-08-24T06:48:52Z

@shen-shanshan can you attach some perf figures ? Thank you.

tjtanaa · 2026-08-24T06:50:02Z

@@ -945,3 +982,22 @@ def wq_b_and_q_quant():
        else:
            q, q_scale = q_quant, None
        return q, q_scale, weights
+
+    def _wq_b_proj(


This function seems to be duplicated.

This function seems to be duplicated.

There are two _wq_b_proj methods (wq_b in SWA and wq_b in indexer). We should skip the quant process in the both linear layers when using this fused kernel.

@tjtanaa I have also attached the profile traces before and after this PR, which could make it easier to understand. 😊

shen-shanshan · 2026-08-24T06:58:02Z

@shen-shanshan can you attach some perf figures ? Thank you.

Yeah, I have updated the benchmark results just now.

mergify · 2026-08-27T01:21:28Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @shen-shanshan.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

shen-shanshan · 2026-08-27T06:20:09Z

/ci run

github-actions · 2026-08-27T06:20:27Z

✅ Triggered Buildkite CI #85760 for commit ec9d70553685.

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

shen-shanshan · 2026-08-27T09:05:52Z

/ci run

github-actions · 2026-08-27T09:06:19Z

✅ Triggered Buildkite CI #85791 for commit 0aad5d851c40.

Rebased onto main after vllm-project#53540 landed the fused_mla_dual_rms_norm_group_quant custom op. The op registration this PR previously carried is dropped; only the accessor the pattern matcher needs is added here. Adds MLADualRMSGroupQuantPattern, the group-quant sibling of the existing MLADualRMSPerTokenQuantPattern, so the MLA FP8 path picks up the fused AITER kernel through the compile pass rather than a hand-wired call site -- covering DeepSeek-R1 / MLA, which vllm-project#53540 does not touch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Andy Friedrich <afriedri@amd.com>

…-V4 (vllm-project#53540) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

…-V4 (vllm-project#53540) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Signed-off-by: mikeshawcode <michaelwshaw2@gmail.com>

…-V4 (vllm-project#53540) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

shen-shanshan requested review from AndreasKaratzas, dllehr-amd, hongxiayang, tjtanaa and zyongye as code owners August 24, 2026 06:46

claude Bot reviewed Aug 24, 2026

View reviewed changes

mergify Bot added deepseek Related to DeepSeek models quantization DSv4 rocm Related to AMD ROCm labels Aug 24, 2026

github-project-automation Bot added this to AMD Aug 24, 2026

github-project-automation Bot moved this to Todo in AMD Aug 24, 2026

tjtanaa reviewed Aug 24, 2026

View reviewed changes

shen-shanshan mentioned this pull request Aug 24, 2026

[RFC]: DeepSeek-V4 Performance Optimization on ROCm (Phase Two) #52911

Open

17 tasks

mergify Bot added the needs-rebase label Aug 27, 2026

Fuse SWA q/kv RMSNorm and q FP8 group quant for DeepSeek-V4

ec9d705

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

shen-shanshan force-pushed the rocm-dsv4-csa-kernel-fuse branch from 479de2f to ec9d705 Compare August 27, 2026 06:19

add ut

0aad5d8

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

mergify Bot removed the needs-rebase label Aug 27, 2026

AndreasKaratzas approved these changes Aug 27, 2026

View reviewed changes

AndreasKaratzas merged commit 32ad140 into vllm-project:main Aug 27, 2026
87 checks passed

github-project-automation Bot moved this from Todo to Done in AMD Aug 27, 2026

afriedri mentioned this pull request Aug 27, 2026

[ROCm][fusion] add dual RMSNorm quant fusion pattern #50317

Open

jiacao-amd mentioned this pull request Sep 1, 2026

[Klaud Cold] [AMD] Bump DSV4 MI355X vLLM 8k/1k to the 2026-09-01 nightly / [Klaud Cold] [AMD] 将 DSV4 MI355X vLLM 8k/1k 更新至 2026-09-01 nightly SemiAnalysisAI/InferenceX#2792

Merged

eky-amd mentioned this pull request Sep 2, 2026

[ROCm] Fuse MLA dual RMSNorm + FP8 group quant for DeepSeek-R1 #54857

Draft

4 tasks

Navigation Menu

Uh oh!

[ROCm][Perf] Fuse SWA q/kv RMSNorm and q FP8 group quant for DeepSeek-V4 - #53540

[ROCm][Perf] Fuse SWA q/kv RMSNorm and q FP8 group quant for DeepSeek-V4#53540
AndreasKaratzas merged 2 commits into
vllm-project:mainfrom
shen-shanshan:rocm-dsv4-csa-kernel-fuse

shen-shanshan commented Aug 24, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

tjtanaa commented Aug 24, 2026

Uh oh!

tjtanaa Aug 24, 2026

Uh oh!

shen-shanshan Aug 24, 2026 •

edited

Loading

Uh oh!

shen-shanshan Aug 24, 2026

Uh oh!

shen-shanshan commented Aug 24, 2026

Uh oh!

mergify Bot commented Aug 27, 2026

Uh oh!

shen-shanshan commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026

Uh oh!

shen-shanshan commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

shen-shanshan commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Benchmark

Accuracy Test

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

tjtanaa commented Aug 24, 2026

Uh oh!

tjtanaa Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

shen-shanshan Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shen-shanshan Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

shen-shanshan commented Aug 24, 2026

Uh oh!

mergify Bot commented Aug 27, 2026

Uh oh!

shen-shanshan commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026

Uh oh!

shen-shanshan commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shen-shanshan commented Aug 24, 2026 •

edited

Loading

shen-shanshan Aug 24, 2026 •

edited

Loading