[Fusion] [Graph] Add qknorm rope fusion operator by wxsIcey · Pull Request #4711 · vllm-project/vllm-ascend

wxsIcey · 2025-12-04T09:41:55Z

What this PR does / why we need it?

This PR add qkv_rmsnorm_rope operator and introduces a graph fusion pass for qknorm_rope operations. The implementation includes a new configuration flag, a pattern matching pass using torch._inductor.pattern_matcher, and a custom Triton kernel for the fused operation.

Co-authored-by: Angazenn supperccell@163.com

Does this PR introduce any user-facing change?

Yes, add new additional_config

How was this patch tested?

local test

vLLM version: v0.12.0
vLLM main: vllm-project/vllm@ad32e3e

gemini-code-assist

Code Review

This pull request introduces a graph fusion pass for qknorm_rope operations on Ascend hardware, which is a great step for performance optimization. The implementation includes a new configuration flag, a pattern matching pass using torch._inductor.pattern_matcher, and a custom Triton kernel for the fused operation. The code is well-structured, but I've identified several areas for improvement regarding code quality, robustness, and maintainability. My review comments focus on removing debug artifacts, improving code clarity and consistency, enhancing robustness by avoiding hardcoded values and unsafe module-level initializations, and addressing significant code duplication.

github-actions · 2025-12-04T10:04:32Z

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

A PR should do only one thing, smaller PRs enable faster reviews.
Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

weijinqian0 · 2025-12-08T10:40:47Z

+    return q_output, k_output, v_output
+
+
+direct_register_custom_op(op_name="qkv_rmsnorm_rope",


use import torch_npu._inductor

The pattern_matcher method of inductor does not support the triton operator. It does support torch.ops.aten (aten operator), torch.ops.npu (custom operator), and torch.add (PyTorch API). Therefore, it is wrapped as a custom op.

whx-sjtu · 2025-12-08T10:44:38Z

+    return driver.active.utils.get_device_properties(device)
+
+
+num_vectorcore = get_npu_properties()["num_vectorcore"]


this parameter has already been defined in triton/utils.py

Thanks. I have modified it.

github-actions · 2025-12-10T01:49:48Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

github-actions · 2025-12-11T03:52:28Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

momo609 · 2025-12-11T07:05:46Z

+
+            return q_rope, k_rope, v
+
+        def replacement(qkv: torch.Tensor, q_weight: torch.Tensor,


pattern in 'if xxx else: torch.ops.vllm.qkv_rmsnorm_rope ’ need support in future releases

We don't perform any special checks in the pattern. You can add a new pattern match.

github-actions · 2025-12-11T09:58:17Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

github-actions · 2025-12-12T09:29:45Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

Signed-off-by: wxsIcey <1790571317@qq.com>

wangxiyuan · 2025-12-17T00:54:04Z

Additional config doc should be updated. Please submit a follow-up PR. @wxsIcey

### What this PR does / why we need it? This PR add `qkv_rmsnorm_rope` operator and introduces a graph fusion pass for `qknorm_rope` operations. The implementation includes a new configuration flag, a pattern matching pass using `torch._inductor.pattern_matcher`, and a custom Triton kernel for the fused operation. Co-authored-by: Angazenn [supperccell@163.com](mailto:supperccell@163.com) ### Does this PR introduce _any_ user-facing change? Yes, add new additional_config - vLLM version: v0.12.0 - vLLM main: vllm-project/vllm@ad32e3e --------- Signed-off-by: wxsIcey <1790571317@qq.com>

### What this PR does / why we need it? This PR add `qkv_rmsnorm_rope` operator and introduces a graph fusion pass for `qknorm_rope` operations. The implementation includes a new configuration flag, a pattern matching pass using `torch._inductor.pattern_matcher`, and a custom Triton kernel for the fused operation. Co-authored-by: Angazenn [supperccell@163.com](mailto:supperccell@163.com) ### Does this PR introduce _any_ user-facing change? Yes, add new additional_config - vLLM version: v0.12.0 - vLLM main: vllm-project/vllm@ad32e3e --------- Signed-off-by: wxsIcey <1790571317@qq.com> Signed-off-by: zrj026 <zhangrunjiang026@gmail.com>

gemini-code-assist Bot reviewed Dec 4, 2025

View reviewed changes

github-actions Bot added module:ops module:core labels Dec 4, 2025

wxsIcey changed the title ~~[Fusion] [Graph] Add qknorm rope fusion~~ [Fusion] [Graph] Add qknorm rope fusion operator Dec 5, 2025

github-actions Bot added the module:tests label Dec 7, 2025

Angazenn reviewed Dec 8, 2025

View reviewed changes

Comment thread vllm_ascend/worker/model_runner_v1.py Outdated

wangxiyuan approved these changes Dec 8, 2025

View reviewed changes

weijinqian0 reviewed Dec 8, 2025

View reviewed changes

whx-sjtu requested changes Dec 8, 2025

View reviewed changes

wxsIcey marked this pull request as ready for review December 9, 2025 02:04

github-actions Bot added merge-conflicts and removed module:tests labels Dec 9, 2025

wxsIcey force-pushed the qknorm_rope_fusion branch from bac8b40 to 59f15a7 Compare December 11, 2025 02:18

github-actions Bot removed the merge-conflicts label Dec 11, 2025

wxsIcey added ready read for review ready-for-test start test by label for PR labels Dec 11, 2025

wxsIcey requested a review from whx-sjtu December 11, 2025 03:02

github-actions Bot added the merge-conflicts label Dec 11, 2025

wxsIcey force-pushed the qknorm_rope_fusion branch from 5a890c9 to 98a4d21 Compare December 11, 2025 04:54

github-actions Bot removed the merge-conflicts label Dec 11, 2025

momo609 reviewed Dec 11, 2025

View reviewed changes

github-actions Bot added the merge-conflicts label Dec 11, 2025

wxsIcey force-pushed the qknorm_rope_fusion branch from 3d8050a to 95718dd Compare December 12, 2025 03:46

github-actions Bot added merge-conflicts and removed merge-conflicts labels Dec 12, 2025

wxsIcey added 21 commits December 16, 2025 14:29

fix triton

fdb549c

Signed-off-by: wxsIcey <1790571317@qq.com>

fix

ead3622

Signed-off-by: wxsIcey <1790571317@qq.com>

install triton

4a02ac5

Signed-off-by: wxsIcey <1790571317@qq.com>

fix ut

9019f16

Signed-off-by: wxsIcey <1790571317@qq.com>

fix

09019ab

Signed-off-by: wxsIcey <1790571317@qq.com>

fix ut

5454d63

Signed-off-by: wxsIcey <1790571317@qq.com>

fix ut

360fcf5

Signed-off-by: wxsIcey <1790571317@qq.com>

resolve conflict

129fcde

Signed-off-by: wxsIcey <1790571317@qq.com>

change workflow

536ead4

Signed-off-by: wxsIcey <1790571317@qq.com>

tiny fix

8599237

Signed-off-by: wxsIcey <1790571317@qq.com>

fix

5de88fc

Signed-off-by: wxsIcey <1790571317@qq.com>

fix

fe8a177

Signed-off-by: wxsIcey <1790571317@qq.com>

fix qwen3 next

2c43712

Signed-off-by: wxsIcey <1790571317@qq.com>

fix eagle

519aa97

Signed-off-by: wxsIcey <1790571317@qq.com>

fix eagle

5e22e20

Signed-off-by: wxsIcey <1790571317@qq.com>

fix mla cp

5e6a838

Signed-off-by: wxsIcey <1790571317@qq.com>

fix

993744a

Signed-off-by: wxsIcey <1790571317@qq.com>

fix e2e

31aa1b2

Signed-off-by: wxsIcey <1790571317@qq.com>

recover e2e

01814e5

Signed-off-by: wxsIcey <1790571317@qq.com>

resolve conflict

9143521

Signed-off-by: wxsIcey <1790571317@qq.com>

fix

dd4617f

Signed-off-by: wxsIcey <1790571317@qq.com>

wxsIcey force-pushed the qknorm_rope_fusion branch from cf04ffc to dd4617f Compare December 16, 2025 14:29

github-actions Bot removed the merge-conflicts label Dec 16, 2025

fix

d90be75

Signed-off-by: wxsIcey <1790571317@qq.com>

wangxiyuan approved these changes Dec 17, 2025

View reviewed changes

wangxiyuan merged commit cadfa5d into vllm-project:main Dec 17, 2025
25 checks passed

		return q_output, k_output, v_output


		direct_register_custom_op(op_name="qkv_rmsnorm_rope",

		return driver.active.utils.get_device_properties(device)


		num_vectorcore = get_npu_properties()["num_vectorcore"]


		return q_rope, k_rope, v

		def replacement(qkv: torch.Tensor, q_weight: torch.Tensor,

Conversation

wxsIcey commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this PR does / why we need it?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Dec 4, 2025

Uh oh!

Uh oh!

weijinqian0 Dec 8, 2025

Choose a reason for hiding this comment

Uh oh!

wxsIcey Dec 9, 2025

Choose a reason for hiding this comment

Uh oh!

whx-sjtu Dec 8, 2025

Choose a reason for hiding this comment

Uh oh!

wxsIcey Dec 9, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Dec 10, 2025

Uh oh!

github-actions Bot commented Dec 11, 2025

Uh oh!

momo609 Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

wxsIcey Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Dec 11, 2025

Uh oh!

github-actions Bot commented Dec 12, 2025

Uh oh!

Uh oh!

wangxiyuan commented Dec 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

wxsIcey commented Dec 4, 2025 •

edited

Loading