[Refactor]Reconstructing TP asymmetry and C8 quantization buffer allocation by pichangping · Pull Request #7658 · vllm-project/vllm-ascend

pichangping · 2026-03-26T02:15:19Z

What this PR does / why we need it?

The legacy issue of pr #7222 ：Unifying and optimizing the creation of buffers for TP asymmetry and C8 quantization in PD separation scenarios.

Does this PR introduce any user-facing change?

no

How was this patch tested?

vLLM version: v0.18.0
vLLM main: vllm-project/vllm@35141a7

…DSV3.1 C8. Signed-off-by: pichangping <1337510399@qq.com>

Signed-off-by: pichangping <1337510399@qq.com>

…o main

Signed-off-by: pichangping <1337510399@qq.com>

github-actions · 2026-03-26T02:15:37Z

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

gemini-code-assist · 2026-03-26T02:15:38Z

Summary of Changes

Hello, 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 addresses a legacy issue by refactoring and unifying the buffer allocation strategy for KV caches, specifically for scenarios involving Tensor Parallelism (TP) asymmetry and C8 quantization. By consolidating dedicated quantization buffers into the main KV buffers, the change simplifies memory management and streamlines the attention processing logic, leading to a more robust and efficient system.

Highlights

KV Cache Buffer Unification: The dedicated k_quant_buffer and v_quant_buffer have been removed, and their functionality is now integrated into the main k_buffer and v_buffer. This unifies buffer allocation for both Tensor Parallelism (TP) asymmetry and C8 quantization.
Attention Logic Streamlining: The attention mechanism's _forward_decode method now uses a unified logic for input layout and q_nope processing when fa_quant_layer or enable_kv_nz is active, removing a redundant elif block.
Dynamic Dequantization Scale Reshaping: Added logic to dynamically reshape dequant_scale_q_nope based on whether speculative decoding (MTP) is enabled or disabled, ensuring correct dimensions for the FIA (Fused Attention) operation.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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.

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

gemini-code-assist

Code Review

This pull request refactors KV cache buffer management by consolidating separate quantization buffers into the main k_buffer and v_buffer. It updates the _forward_decode and _mla_preprocess_only_decode functions to correctly handle fa_quant_layer and dequant_scale_q_nope reshaping, integrating quantization logic more directly into the main KV buffer handling. A critical issue was identified where the v_buffer might be initialized with an incorrect data type, potentially causing mismatches in mixed-precision quantization scenarios.

gemini-code-assist · 2026-03-26T02:18:45Z

-            )
-            self.k_quant_buffer = align_memory(self.k_quant_buffer, alignment)[:quant_k_cache_numel].view(
-                -1, first_kv_cache.shape[-1]
+                first_v_cache_numel + alignment, dtype=first_k_cache.dtype, device=first_k_cache.device


The v_buffer is being created with first_k_cache.dtype. In scenarios with mixed-precision KV cache quantization (like C8 where K is int8 and V is float16/bfloat16), this could lead to a dtype mismatch. It seems first_v_cache.dtype should be used here to ensure the buffer for V cache has the correct data type, especially when self.enable_kv_quant is true.

Suggested change

first_v_cache_numel + alignment, dtype=first_k_cache.dtype, device=first_k_cache.device

first_v_cache_numel + alignment, dtype=first_v_cache.dtype, device=first_v_cache.device

kvcache should reside on a single device, so the device here is always named first_k_cache.device.

Signed-off-by: pichangping <1337510399@qq.com>

kunpengW-code · 2026-03-26T08:30:57Z

+            """When MTP is enabled or disabled, the different input_layout results in different
+            dimensions of dequant_scale_q_nope required by FIA.
+            """
+            if self.speculative_config is None:


The attn_metadata.attn_state status is not checked. However, since C8 quantization is performed only at the D node, the condition on line 1300~1305 is already met, so there is no problem.

github-actions · 2026-03-31T14:19:24Z

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

…o main # Conflicts: # vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_layerwise_connector.py

github-actions · 2026-04-21T08:47:20Z

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

pichangping added 9 commits March 23, 2026 21:47

bugfix: Fixed the error issue when overlaying MTP and full decode on …

53a587a

…DSV3.1 C8. Signed-off-by: pichangping <1337510399@qq.com>

fix CI

d5a1ac2

Signed-off-by: pichangping <1337510399@qq.com>

fix CI

fa9d361

Signed-off-by: pichangping <1337510399@qq.com>

Merge branch 'main' into main

f1e482f

Merge branch 'vllm-project:main' into main

4351728

Refactoring ds C8

4560f93

Signed-off-by: pichangping <1337510399@qq.com>

Merge branch 'main' of https://github.com/pichangping/vllm-ascend int…

051baa7

…o main

Merge branch 'vllm-project:main' into main

de3decc

Reconstructing TP asymmetry and C8 quantization buffer allocation

84ad5c0

Signed-off-by: pichangping <1337510399@qq.com>

pichangping requested review from LCAIZJ, MengqingCao, weijinqian0 and whx-sjtu as code owners March 26, 2026 02:15

gemini-code-assist bot reviewed Mar 26, 2026

View reviewed changes

pichangping added 2 commits March 26, 2026 10:19

fix ci

7144bfb

Signed-off-by: pichangping <1337510399@qq.com>

fix ci

7b8b60a

Signed-off-by: pichangping <1337510399@qq.com>

pichangping changed the title ~~[refector]Reconstructing TP asymmetry and C8 quantization buffer allocation~~ [refactor]Reconstructing TP asymmetry and C8 quantization buffer allocation Mar 26, 2026

pichangping and others added 3 commits March 26, 2026 11:40

Remove duplicate code

244fb7b

Signed-off-by: pichangping <1337510399@qq.com>

fix ci

cd86289

Signed-off-by: pichangping <1337510399@qq.com>

Merge branch 'main' into main

f3d464c

wangxiyuan added ready read for review ready-for-test start test by label for PR labels Mar 26, 2026

kunpengW-code reviewed Mar 26, 2026

View reviewed changes

pichangping added 2 commits March 26, 2026 19:03

Merge branch 'main' into main

252c6e4

Merge branch 'main' into main

31885f6

weijinqian0 approved these changes Mar 27, 2026

View reviewed changes

github-actions bot added the merge-conflicts label Mar 31, 2026

pichangping closed this Apr 3, 2026

pichangping force-pushed the main branch from 31885f6 to d39031b Compare April 3, 2026 02:40

Merge branch 'main' of https://github.com/pichangping/vllm-ascend int…

207fecd

…o main # Conflicts: # vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_layerwise_connector.py

pichangping reopened this Apr 3, 2026

github-actions bot removed the merge-conflicts label Apr 3, 2026

pichangping changed the title ~~[refactor]Reconstructing TP asymmetry and C8 quantization buffer allocation~~ [Refactor]Reconstructing TP asymmetry and C8 quantization buffer allocation Apr 3, 2026

github-actions bot added the merge-conflicts label Apr 21, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Refactor]Reconstructing TP asymmetry and C8 quantization buffer allocation#7658

[Refactor]Reconstructing TP asymmetry and C8 quantization buffer allocation#7658
pichangping wants to merge 17 commits intovllm-project:mainfrom
pichangping:main

pichangping commented Mar 26, 2026 •

edited

Loading

Uh oh!

github-actions bot commented Mar 26, 2026

Uh oh!

gemini-code-assist bot commented Mar 26, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Mar 26, 2026

Uh oh!

pichangping Mar 26, 2026

Uh oh!

kunpengW-code Mar 26, 2026

Uh oh!

pichangping Mar 26, 2026

Uh oh!

github-actions bot commented Mar 31, 2026

Uh oh!

github-actions bot commented Apr 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	first_v_cache_numel + alignment, dtype=first_k_cache.dtype, device=first_k_cache.device
	first_v_cache_numel + alignment, dtype=first_v_cache.dtype, device=first_v_cache.device

Conversation

pichangping commented Mar 26, 2026 • 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!

github-actions bot commented Mar 26, 2026

Uh oh!

gemini-code-assist bot commented Mar 26, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

pichangping Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

kunpengW-code Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

pichangping Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Mar 31, 2026

Uh oh!

github-actions bot commented Apr 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pichangping commented Mar 26, 2026 •

edited

Loading