[bugfix]Fix accuracy issue in PCP/DCP with speculative decoding

[kunpengW-code]

What this PR does / why we need it?

This PR fixes an accuracy issue that occurs when using Prefill/Decode Context Parallelism (PCP/DCP) in conjunction with speculative decoding (MTP). The issue is caused by an irregular attention mask shape when both features are enabled.

The fix involves flattening the block_table for speculative decoding requests under PCP/DCP to ensure a regular attention mask. This PR also introduces a use_cp property for cleaner code and updates dummy runs to handle this scenario correctly.

Does this PR introduce any user-facing change?

No. This is a bug fix that improves accuracy and should not have user-facing API changes.

How was this patch tested?

• vLLM version: v0.15.0
• vLLM main: https://github.com/vllm-project/vllm/commit/v0.15.0

[gemini-code-assist]

Summary of Changes

Hello @kunpengW-code, 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 an accuracy problem that arises when Parallel Context Processing (PCP) or Data Parallel Processing (DCP) is combined with Multi-Token Prediction (MTP) or speculative decoding. The changes streamline the detection of parallel processing through a dedicated flag and implement crucial adjustments to how block_table and query_lens are managed under these specific conditions. These modifications are designed to prevent irregular attention mask shapes and ensure accurate token handling during speculative decoding with context parallelism.

Highlights

• Refactored Parallel Processing Check: The condition for checking if Parallel Context Processing (PCP) or Data Parallel Processing (DCP) is active has been refactored from self.pcp_size * self.dcp_size > 1 to a more concise self.use_cp flag across several functions.
• Block Table Flattening for Speculative Decoding: New logic has been introduced to flatten the block_table when both speculative decoding and context parallelism are active. This change specifically addresses potential issues with irregular attention mask shapes by adjusting block_table_gid_0 and query_lens.
• Query Lengths Management for Speculative Decoding: The _dummy_run function now includes an update to self.pcp_manager.query_lens_pcp_full for speculative decoding scenarios when context parallelism is enabled, ensuring correct token length tracking.

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

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. 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 addresses an accuracy issue with Prefill/Decode Context Parallelism (PCP/DCP) when used with speculative decoding (MTP). The changes primarily involve flattening the block table to prevent irregular attention mask shapes, which was the root cause of the problem. The PR also introduces a use_cp property to simplify checking for context parallelism, improving code readability. The changes appear correct and well-targeted to fix the described issue.

As per the repository's style guide, I've suggested an updated Pull Request title and summary below:

Suggested PR Title:

[Worker][BugFix] Fix accuracy issue in PCP/DCP with speculative decoding

Suggested PR Summary:

### What this PR does / why we need it?

This PR fixes an accuracy issue that occurs when using Prefill/Decode Context Parallelism (PCP/DCP) in conjunction with speculative decoding (MTP). The issue is caused by an irregular attention mask shape when both features are enabled.

The fix involves flattening the `block_table` for speculative decoding requests under PCP/DCP to ensure a regular attention mask. This PR also introduces a `use_cp` property for cleaner code and updates dummy runs to handle this scenario correctly.

### Does this PR introduce _any_ user-facing change?

No. This is a bug fix that improves accuracy and should not have user-facing API changes.

### How was this patch tested?

CI passed with new added/existing test.

[github-actions]

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

[pisceskkk]

Maybe we should extract this code snippet into a function of PCPManager?

[kunpengW-code]

Yes, this is more appropriate; we can integrate this logic with the generate_pcp_metadata method.

[kunpengW-code]

deepseekv3.1 w8a8 tp8 pcp2 dcp8 mtp1
gsm8k accuracy:

dataset	version	metric	mode	vllm-api-general
gsm8kdataset	-	accuracy	gen	96.29