[feat]: oproj tensor parallelism in pure DP and graph-mode scenarios. #2167

lidenghui1110 · 2025-08-01T07:03:52Z

What this PR does / why we need it?

This PR introduces Oproj matrix tensor model parallel to achieve decreasing of memory consumption. It only support graph mode in pure DP scenario.

In deepseek r1 w8a8 PD disagregated Decode instance, using pure DP, with oproj_tensor_parallel_size = 8, we have 1 ms TPOT increasing, saved 5.8 GB NPU memory per RANK. We got best performance when oproj_tensor_parallel_size=4 without TPOT increasing.

performance data:

Does this PR introduce any user-facing change?

This PR introduces one new config in additional_config.

Name	Effect	Required	Type	Constraints
oproj_tensor_parallel_size	Split the o_proj matrix along the row dimension (head num * head dim) into oproj_tensor_parallel_size pieces.	No	int	default value is None, once this value is set, the feature will be enabled, head num * head dim must be divisible by this value.

example

--additional_config={"oproj_tensor_parallel_size": 8}

How was this patch tested?

vLLM version: v0.10.1.1
vLLM main: vllm-project/vllm@eddaafc

github-actions · 2025-08-01T07:38:40Z

👋 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.

sdmyzlp · 2025-08-11T01:36:01Z

vllm_ascend/ops/linear.py

+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):


This function seems to be identical with that of RowParallelLinear, why do we need to rewrite it here?

in origin weight_load,

tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size()

we need replace it into custom comm group

tp_rank = self.tp_rank tp_size = self.tp_size

It seems that the latest vllm does not have this problem.

Understood, thanks

github-actions · 2025-08-11T07:39:05Z

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

github-actions · 2025-08-11T08:03:32Z

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

github-actions · 2025-08-12T06:14:29Z

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

ApsarasX · 2025-08-12T06:50:43Z

vllm_ascend/quantization/quant_config.py

+            else:
+                tp_rank = get_tensor_model_parallel_rank()
+        else:
+            tp_rank = 0


What means tp_rank = 0?

This origin code here

if isinstance(layer, RowParallelLinear): tp_rank = get_tensor_model_parallel_rank() return self.quant_method.apply(layer, x, bias, tp_rank) return self.quant_method.apply(layer, x, bias)

The default situation is not passing tp, which is tp=0

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

zzhx1 · 2025-09-06T16:00:59Z

@wangxiyuan This PR is ready，and also fixed the bug related to linearBase.

wangxiyuan · 2025-09-07T02:30:13Z

vllm_ascend/patch/worker/patch_common/patch_lora.py

@@ -0,0 +1,15 @@
+import vllm


looks that these 3 file can merged into one

…vllm-project#2167) ### What this PR does / why we need it? This PR introduces Oproj matrix tensor model parallel to achieve decreasing of memory consumption. It only support graph mode in pure DP scenario. In deepseek r1 w8a8 PD disagregated Decode instance, using pure DP, with oproj_tensor_parallel_size = 8, we have 1 ms TPOT increasing, saved 5.8 GB NPU memory per RANK. We got best performance when oproj_tensor_parallel_size=4 without TPOT increasing. performance data: <img width="1442" height="442" alt="image" src="https://github.com/user-attachments/assets/83270fc5-868a-4387-b0a9-fac29b4a376d" /> ### Does this PR introduce _any_ user-facing change? This PR introduces one new config in `additional_config`. | Name | Effect | Required | Type | Constraints | | :---------------------------- | :--------------------------------------- | :------- | :--- | :----------------- | | oproj_tensor_parallel_size | Split the o_proj matrix along the row dimension (head num * head dim) into oproj_tensor_parallel_size pieces. | No | int | default value is None, once this value is set, the feature will be enabled, head num * head dim must be divisible by this value. | example `--additional_config={"oproj_tensor_parallel_size": 8}` ### How was this patch tested? - vLLM version: v0.10.1.1 - vLLM main: vllm-project/vllm@eddaafc --------- Signed-off-by: zzhx1 <zzh_201018@outlook.com> Co-authored-by: zzh <zzh_201018@outlook.com>

…vllm-project#2167) ### What this PR does / why we need it? This PR introduces Oproj matrix tensor model parallel to achieve decreasing of memory consumption. It only support graph mode in pure DP scenario. In deepseek r1 w8a8 PD disagregated Decode instance, using pure DP, with oproj_tensor_parallel_size = 8, we have 1 ms TPOT increasing, saved 5.8 GB NPU memory per RANK. We got best performance when oproj_tensor_parallel_size=4 without TPOT increasing. performance data: <img width="1442" height="442" alt="image" src="https://github.com/user-attachments/assets/83270fc5-868a-4387-b0a9-fac29b4a376d" /> ### Does this PR introduce _any_ user-facing change? This PR introduces one new config in `additional_config`. | Name | Effect | Required | Type | Constraints | | :---------------------------- | :--------------------------------------- | :------- | :--- | :----------------- | | oproj_tensor_parallel_size | Split the o_proj matrix along the row dimension (head num * head dim) into oproj_tensor_parallel_size pieces. | No | int | default value is None, once this value is set, the feature will be enabled, head num * head dim must be divisible by this value. | example `--additional_config={"oproj_tensor_parallel_size": 8}` ### How was this patch tested? - vLLM version: v0.10.1.1 - vLLM main: vllm-project/vllm@eddaafc --------- Signed-off-by: zzhx1 <zzh_201018@outlook.com> Co-authored-by: zzh <zzh_201018@outlook.com> Signed-off-by: offline0806 <z00858301@china.huawei.com>

…vllm-project#2167) ### What this PR does / why we need it? This PR introduces Oproj matrix tensor model parallel to achieve decreasing of memory consumption. It only support graph mode in pure DP scenario. In deepseek r1 w8a8 PD disagregated Decode instance, using pure DP, with oproj_tensor_parallel_size = 8, we have 1 ms TPOT increasing, saved 5.8 GB NPU memory per RANK. We got best performance when oproj_tensor_parallel_size=4 without TPOT increasing. performance data: <img width="1442" height="442" alt="image" src="https://github.com/user-attachments/assets/83270fc5-868a-4387-b0a9-fac29b4a376d" /> ### Does this PR introduce _any_ user-facing change? This PR introduces one new config in `additional_config`. | Name | Effect | Required | Type | Constraints | | :---------------------------- | :--------------------------------------- | :------- | :--- | :----------------- | | oproj_tensor_parallel_size | Split the o_proj matrix along the row dimension (head num * head dim) into oproj_tensor_parallel_size pieces. | No | int | default value is None, once this value is set, the feature will be enabled, head num * head dim must be divisible by this value. | example `--additional_config={"oproj_tensor_parallel_size": 8}` ### How was this patch tested? - vLLM version: v0.10.1.1 - vLLM main: vllm-project/vllm@eddaafc --------- Signed-off-by: zzhx1 <zzh_201018@outlook.com> Co-authored-by: zzh <zzh_201018@outlook.com>

…vllm-project#2167) This PR introduces Oproj matrix tensor model parallel to achieve decreasing of memory consumption. It only support graph mode in pure DP scenario. In deepseek r1 w8a8 PD disagregated Decode instance, using pure DP, with oproj_tensor_parallel_size = 8, we have 1 ms TPOT increasing, saved 5.8 GB NPU memory per RANK. We got best performance when oproj_tensor_parallel_size=4 without TPOT increasing. performance data: <img width="1442" height="442" alt="image" src="https://github.com/user-attachments/assets/83270fc5-868a-4387-b0a9-fac29b4a376d" /> This PR introduces one new config in `additional_config`. | Name | Effect | Required | Type | Constraints | | :---------------------------- | :--------------------------------------- | :------- | :--- | :----------------- | | oproj_tensor_parallel_size | Split the o_proj matrix along the row dimension (head num * head dim) into oproj_tensor_parallel_size pieces. | No | int | default value is None, once this value is set, the feature will be enabled, head num * head dim must be divisible by this value. | example `--additional_config={"oproj_tensor_parallel_size": 8}` - vLLM version: v0.10.1.1 - vLLM main: vllm-project/vllm@eddaafc --------- Signed-off-by: zzhx1 <zzh_201018@outlook.com> Co-authored-by: zzh <zzh_201018@outlook.com>

github-actions bot added module:ops module:core module:quantization labels Aug 1, 2025

sdmyzlp reviewed Aug 11, 2025

View reviewed changes

zzhx1 force-pushed the oproj branch from 974d272 to a02e6fd Compare August 11, 2025 07:38

github-actions bot added documentation Improvements or additions to documentation ci/build module:tests module:tools merge-conflicts labels Aug 11, 2025

zzhx1 force-pushed the oproj branch from a02e6fd to 974d272 Compare August 11, 2025 07:45

github-actions bot removed documentation Improvements or additions to documentation ci/build module:tests module:tools merge-conflicts labels Aug 11, 2025

zzhx1 force-pushed the oproj branch 2 times, most recently from 71c3e49 to b925b4b Compare August 11, 2025 08:02

github-actions bot added the merge-conflicts label Aug 11, 2025

zzhx1 force-pushed the oproj branch from b925b4b to 71c3e49 Compare August 11, 2025 08:04

github-actions bot removed the merge-conflicts label Aug 11, 2025

zzhx1 force-pushed the oproj branch from 71c3e49 to e11022b Compare August 11, 2025 08:58

github-actions bot added the merge-conflicts label Aug 12, 2025

github-actions bot added the module:tests label Aug 12, 2025

ApsarasX reviewed Aug 12, 2025

View reviewed changes

zzhx1 force-pushed the oproj branch from 0abb2ad to 506b7d6 Compare August 18, 2025 15:49

zzhx1 force-pushed the oproj branch 8 times, most recently from 51d1def to 1e76d68 Compare September 4, 2025 12:39

zzhx1 added 5 commits September 5, 2025 11:25

feat: oproj tensor parallelism in pure DP and graph-mode scenarios.

365e7e5

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

Improve the overall code related to oproj_tensor_parallel.

c97eec3

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

[CI] fix

10c110e

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

remove patch_linear and integrate matmul allreduce in linear.py together

95c9b10

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

add UT

0a8017a

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

zzhx1 force-pushed the oproj branch 2 times, most recently from 065be1d to 6baf96b Compare September 5, 2025 06:46

fix lora patch

18b14e7

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

zzhx1 force-pushed the oproj branch 3 times, most recently from dcb84e0 to e7aaaf8 Compare September 6, 2025 02:43

fix ut bug

dc8a0c8

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

zzhx1 force-pushed the oproj branch from d8759f2 to a345121 Compare September 6, 2025 10:04

fix linearBase UT bug from vllm

a675b1c

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

zzhx1 force-pushed the oproj branch from c7020fc to a675b1c Compare September 6, 2025 12:23

wangxiyuan approved these changes Sep 7, 2025

View reviewed changes

vllm_ascend/patch/worker/patch_common/patch_lora.py

@@ -0,0 +1,15 @@

import vllm

Copy link

Collaborator

wangxiyuan Sep 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks that these 3 file can merged into one

wangxiyuan merged commit 5a71815 into vllm-project:main Sep 7, 2025
25 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[feat]: oproj tensor parallelism in pure DP and graph-mode scenarios. #2167

[feat]: oproj tensor parallelism in pure DP and graph-mode scenarios. #2167

Uh oh!

lidenghui1110 commented Aug 1, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Aug 1, 2025

Uh oh!

sdmyzlp Aug 11, 2025

Uh oh!

zzhx1 Aug 11, 2025

Uh oh!

sdmyzlp Aug 11, 2025

Uh oh!

github-actions bot commented Aug 11, 2025

Uh oh!

github-actions bot commented Aug 11, 2025

Uh oh!

github-actions bot commented Aug 12, 2025

Uh oh!

ApsarasX Aug 12, 2025

Uh oh!

zzhx1 Aug 12, 2025

Uh oh!

zzhx1 commented Sep 6, 2025

Uh oh!

wangxiyuan Sep 7, 2025

Uh oh!

Uh oh!

Uh oh!

[feat]: oproj tensor parallelism in pure DP and graph-mode scenarios. #2167

[feat]: oproj tensor parallelism in pure DP and graph-mode scenarios. #2167

Uh oh!

Conversation

lidenghui1110 commented Aug 1, 2025 • edited by github-actions bot 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!

github-actions bot commented Aug 1, 2025

Uh oh!

sdmyzlp Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

zzhx1 Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

sdmyzlp Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Aug 11, 2025

Uh oh!

github-actions bot commented Aug 11, 2025

Uh oh!

github-actions bot commented Aug 12, 2025

Uh oh!

ApsarasX Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

zzhx1 Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

zzhx1 commented Sep 6, 2025

Uh oh!

wangxiyuan Sep 7, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

lidenghui1110 commented Aug 1, 2025 •

edited by github-actions bot

Loading