[BugFix] Fix Qwen3.5 MoE flash comm v1 shared expert shape error of mtp layer on A2 by wangbj127 · Pull Request #7683 · vllm-project/vllm-ascend

wangbj127 · 2026-03-26T08:53:39Z

What this PR does / why we need it?

Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled.
fix #7996

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

gemini-code-assist · 2026-03-26T08:53:55Z

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 specific bugs affecting the Qwen3.5 series models when running on the A2 platform. The changes primarily refine the conditional logic for tensor communication and reduction operations, ensuring that multimodal and vision-language models correctly interact with flash_comm_v1 optimizations and distributed communication primitives under various operational modes.

Highlights

Conditional Flash Communication Bypass: Adjusted the condition for bypassing flash_comm_v1_enabled in _maybe_pad_and_reduce_impl to correctly handle cases where a draft Vision-Language (VL) model is not involved in endpoint communication, ensuring proper tensor reduction behavior.
Multimodal Model Tensor Reduction Logic: Modified the maybe_pad_and_reduce function in eagle_proposer.py to ensure that multimodal models correctly utilize flash_comm_v1_enabled for tensor operations, resolving an issue where they might have incorrectly bypassed these operations within the 'mtp' method.

🧠 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 refines the conditions for applying tensor parallel communication and padding/reduction operations within the vllm_ascend framework. It modifies vllm_ascend/ops/register_custom_ops.py to include is_ep_comm in the bypass condition for draft and VL models when flash_comm_v1 is disabled. Additionally, in vllm_ascend/spec_decode/eagle_proposer.py, it adjusts the maybe_pad_and_reduce function to prevent multimodal models from undergoing padding and reduction when flash_comm_v1 is enabled, by removing an early return and adding a not self.is_multimodal_model check to the reduction logic. These changes aim to ensure correct behavior and potentially optimize performance for multimodal and other specific model configurations under different communication settings. There are no review comments to address.

github-actions · 2026-03-26T09:21:44Z

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

Signed-off-by: wangbj127 <wangbj1207@126.com>

github-actions · 2026-04-01T14:19:51Z

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

Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

github-actions · 2026-04-08T08:13:11Z

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

Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com> Signed-off-by: anning-2026 <anning@isrc.iscas.ac.cn>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com> Signed-off-by: guxin108 <1252896542@qq.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com> Signed-off-by: zouyida2052 <zouyida2002@gmail.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com> Signed-off-by: nanxing <1014662416@qq.com>

…tp layer on A2 (vllm-project#7683) ### What this PR does / why we need it? Fix Qwen3.5 MoE MTP layer shared expert shape error when flash comm v1 is enabled. fix vllm-project#7996 ### 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 Signed-off-by: wangbj127 <wangbj1207@126.com> Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

Remove the legacy draft VL TP all-reduce bypass from the EP sequence-parallel finalize path and add regression coverage. ### What this PR does / why we need it? This PR fixes a token-layout mismatch when Qwen3.5 MTP is used together with expert parallelism and upstream MoE sequence parallelism. With upstream MoE sequence parallelism enabled, `PrepareAndFinalizeWithAllGather` uses the following communication flow: ```text TP AllGather -> Attention -> TP ReduceScatter -> EP AllGather -> MoE -> EP ReduceScatter ``` Therefore, `maybe_pad_and_reduce` is an EP sequence-parallel finalize operation and must return the local EP/SP sequence shard. However, it still contained a legacy draft-VL special case: ```python if _EXTRA_CTX.is_draft_model and is_vl_model(): return tensor_model_parallel_all_reduce(x) ``` TP AllReduce preserves the global-token layout instead of restoring the local EP/SP shard. For Qwen3.5 MTP, this causes the draft model to return a tensor with an unexpected token dimension, resulting in a deterministic shape mismatch during profiling or execution, for example: ```text 4096 vs 1024 ``` The real and fake implementations of this custom operation were also inconsistent in this case: - The real implementation returned the global-token shape through TP AllReduce. - The fake implementation always returned the local EP shard shape. This special case originated from the FlashComm1 implementation. PR #7683 explicitly restricted the draft-VL TP AllReduce bypass to non-EP communication using `not is_ep_comm`. After FlashComm1 was removed in #13946, `maybe_pad_and_reduce` became an EP-only operation, but the legacy bypass remained without that guard. This PR: - Removes the obsolete draft-VL TP AllReduce bypass. - Removes the corresponding unused imports. - Makes the real implementation follow the same EP ReduceScatter contract as the fake implementation. - Adds a regression UT ensuring that draft-VL execution cannot bypass EP ReduceScatter. - Updates the existing four-card Qwen3.5 MTP test to explicitly cover: - DP=2 - TP=2 - Expert parallelism - Upstream MoE sequence parallelism - MTP with 3 speculative tokens The fix is intentionally limited to the EP/SP finalize communication layer. It does not add model-specific logic or introduce another sequence-parallel configuration flag. Related PRs: #7683, #13946 ### Does this PR introduce _any_ user-facing change? Yes. Qwen3.5 MoE models can now use MTP speculative decoding together with DP, TP, EP, and upstream MoE sequence parallelism without failing because of a token-dimension mismatch. There are no API or configuration interface changes. ### How was this patch tested? Regression coverage added: - Added `test_sp_ep_reduce_scatter_draft_vl_keeps_ep_layout`. - The test reproduces the historical draft-VL state and fails if the EP/SP finalize path attempts to use TP AllReduce. - The expected output remains the current EP rank's local token shard. - Updated the existing Qwen3.5 four-card E2E case to explicitly use: - `Qwen3.5-35B-A3B` - DP=2 - TP=2 - Expert parallelism enabled - `allgather_reducescatter` - Qwen3.5 MTP with 3 speculative tokens - `FULL_DECODE_ONLY` Local static validation: - Python syntax compilation passed for all three modified files. - `git diff --check` passed. - The final diff contains only the custom operation, regression UT, and Qwen3.5 E2E test changes. Hardware root-cause validation was performed on the equivalent fix: | Configuration | Result | | --- | --- | | SP disabled + MTP enabled | Passed | | SP enabled + MTP disabled | Passed | | SP enabled + MTP enabled with the original bypass | Failed | | SP enabled + MTP enabled after removing the draft-VL TP AllReduce bypass | Passed | Eager-mode shape tracing confirmed: ```text Before finalize: (4096, 2048) Legacy TP AllReduce result: (4096, 2048) EP ReduceScatter result: (1024, 2048) ``` The complete repository UT and the updated four-card E2E test were not executed in the local Windows workspace because it does not contain the required vLLM, PyTorch, torch-npu, pytest, or NPU runtime environment. They are left for the project CI and NPU test environment. - vLLM main: vllm-project/vllm@84030bb --------- Signed-off-by: muziyuhui666 <lijianfu9@huawei.com>

…ject#16081) Remove the legacy draft VL TP all-reduce bypass from the EP sequence-parallel finalize path and add regression coverage. ### What this PR does / why we need it? This PR fixes a token-layout mismatch when Qwen3.5 MTP is used together with expert parallelism and upstream MoE sequence parallelism. With upstream MoE sequence parallelism enabled, `PrepareAndFinalizeWithAllGather` uses the following communication flow: ```text TP AllGather -> Attention -> TP ReduceScatter -> EP AllGather -> MoE -> EP ReduceScatter ``` Therefore, `maybe_pad_and_reduce` is an EP sequence-parallel finalize operation and must return the local EP/SP sequence shard. However, it still contained a legacy draft-VL special case: ```python if _EXTRA_CTX.is_draft_model and is_vl_model(): return tensor_model_parallel_all_reduce(x) ``` TP AllReduce preserves the global-token layout instead of restoring the local EP/SP shard. For Qwen3.5 MTP, this causes the draft model to return a tensor with an unexpected token dimension, resulting in a deterministic shape mismatch during profiling or execution, for example: ```text 4096 vs 1024 ``` The real and fake implementations of this custom operation were also inconsistent in this case: - The real implementation returned the global-token shape through TP AllReduce. - The fake implementation always returned the local EP shard shape. This special case originated from the FlashComm1 implementation. PR vllm-project#7683 explicitly restricted the draft-VL TP AllReduce bypass to non-EP communication using `not is_ep_comm`. After FlashComm1 was removed in vllm-project#13946, `maybe_pad_and_reduce` became an EP-only operation, but the legacy bypass remained without that guard. This PR: - Removes the obsolete draft-VL TP AllReduce bypass. - Removes the corresponding unused imports. - Makes the real implementation follow the same EP ReduceScatter contract as the fake implementation. - Adds a regression UT ensuring that draft-VL execution cannot bypass EP ReduceScatter. - Updates the existing four-card Qwen3.5 MTP test to explicitly cover: - DP=2 - TP=2 - Expert parallelism - Upstream MoE sequence parallelism - MTP with 3 speculative tokens The fix is intentionally limited to the EP/SP finalize communication layer. It does not add model-specific logic or introduce another sequence-parallel configuration flag. Related PRs: vllm-project#7683, vllm-project#13946 ### Does this PR introduce _any_ user-facing change? Yes. Qwen3.5 MoE models can now use MTP speculative decoding together with DP, TP, EP, and upstream MoE sequence parallelism without failing because of a token-dimension mismatch. There are no API or configuration interface changes. ### How was this patch tested? Regression coverage added: - Added `test_sp_ep_reduce_scatter_draft_vl_keeps_ep_layout`. - The test reproduces the historical draft-VL state and fails if the EP/SP finalize path attempts to use TP AllReduce. - The expected output remains the current EP rank's local token shard. - Updated the existing Qwen3.5 four-card E2E case to explicitly use: - `Qwen3.5-35B-A3B` - DP=2 - TP=2 - Expert parallelism enabled - `allgather_reducescatter` - Qwen3.5 MTP with 3 speculative tokens - `FULL_DECODE_ONLY` Local static validation: - Python syntax compilation passed for all three modified files. - `git diff --check` passed. - The final diff contains only the custom operation, regression UT, and Qwen3.5 E2E test changes. Hardware root-cause validation was performed on the equivalent fix: | Configuration | Result | | --- | --- | | SP disabled + MTP enabled | Passed | | SP enabled + MTP disabled | Passed | | SP enabled + MTP enabled with the original bypass | Failed | | SP enabled + MTP enabled after removing the draft-VL TP AllReduce bypass | Passed | Eager-mode shape tracing confirmed: ```text Before finalize: (4096, 2048) Legacy TP AllReduce result: (4096, 2048) EP ReduceScatter result: (1024, 2048) ``` The complete repository UT and the updated four-card E2E test were not executed in the local Windows workspace because it does not contain the required vLLM, PyTorch, torch-npu, pytest, or NPU runtime environment. They are left for the project CI and NPU test environment. - vLLM main: vllm-project/vllm@84030bb --------- Signed-off-by: muziyuhui666 <lijianfu9@huawei.com>

wangbj127 requested review from realliujiaxu, wangxiyuan, whx-sjtu and zzzzwwjj as code owners March 26, 2026 08:53

gemini-code-assist Bot reviewed Mar 26, 2026

View reviewed changes

github-actions Bot added the module:ops label Mar 26, 2026

wangbj127 changed the title ~~[BugFix] Fix bugs of Qwen3.5 series on A2 platform~~ [BugFix] Fix flash comm v1 bugs of Qwen3.5 series on A2 platform Mar 26, 2026

wangbj127 force-pushed the main branch 2 times, most recently from f7303bd to 5839062 Compare March 26, 2026 13:34

Fix flash comm v1 problems of Qwen3.5 series on A2

169ef76

Signed-off-by: wangbj127 <wangbj1207@126.com>

wangbj127 force-pushed the main branch from 5839062 to 169ef76 Compare March 27, 2026 03:25

github-actions Bot added the merge-conflicts label Apr 1, 2026

Merge branch 'main' into main

52544a6

Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

github-actions Bot removed the merge-conflicts label Apr 2, 2026

wangbj127 changed the title ~~[BugFix] Fix flash comm v1 bugs of Qwen3.5 series on A2 platform~~ [BugFix] Fix flash comm v1 bugs of Qwen3.5 series Apr 3, 2026

Merge branch 'main' into main

609803c

wangbj127 changed the title ~~[BugFix] Fix flash comm v1 bugs of Qwen3.5 series~~ [BugFix] Fix flash comm v1 bug of Qwen3.5 MoE on A2 Apr 7, 2026

shaopeng-666 mentioned this pull request Apr 7, 2026

[Bug]: MoE + MTP + FLASHCOMM1 causes tensor dimension mismatch in Qwen3.5-35B #7996

Closed

wangbj127 changed the title ~~[BugFix] Fix flash comm v1 bug of Qwen3.5 MoE on A2~~ [BugFix] Fix Qwen3.5 MoE flash comm v1 shared expert shape error on A2 Apr 7, 2026

wangbj127 changed the title ~~[BugFix] Fix Qwen3.5 MoE flash comm v1 shared expert shape error on A2~~ [BugFix] Fix Qwen3.5 MoE flash comm v1 shared expert shape error of mtp layer on A2 Apr 7, 2026

Angazenn added ready labels Apr 7, 2026

Merge branch 'main' into main

967bfb5

github-actions Bot added the merge-conflicts label Apr 8, 2026

Merge branch 'main' into main

6a601d3

Signed-off-by: wangbj127 <256472688+wangbj127@users.noreply.github.com>

github-actions Bot removed the merge-conflicts label Apr 13, 2026

wangxiyuan approved these changes Apr 13, 2026

View reviewed changes

wangxiyuan merged commit 880e20f into vllm-project:main Apr 13, 2026
51 checks passed

muziyuhui666 mentioned this pull request Sep 8, 2026

[BugFix][Core] Fix MTP MoE finalize layout with upstream SP #16081

Merged

Navigation Menu

[BugFix] Fix Qwen3.5 MoE flash comm v1 shared expert shape error of mtp layer on A2 - #7683

[BugFix] Fix Qwen3.5 MoE flash comm v1 shared expert shape error of mtp layer on A2#7683
wangxiyuan merged 5 commits into
vllm-project:mainfrom
wangbj127:main

wangbj127 commented Mar 26, 2026 •

edited by MengqingCao

Loading

Uh oh!

gemini-code-assist Bot commented Mar 26, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

github-actions Bot commented Mar 26, 2026

Uh oh!

github-actions Bot commented Apr 1, 2026

Uh oh!

github-actions Bot commented Apr 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

wangbj127 commented Mar 26, 2026 • edited by MengqingCao 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 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!

github-actions Bot commented Mar 26, 2026

Uh oh!

github-actions Bot commented Apr 1, 2026

Uh oh!

github-actions Bot commented Apr 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wangbj127 commented Mar 26, 2026 •

edited by MengqingCao

Loading