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

[muziyuhui666]

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:

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:

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:

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:

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

[gemini-code-assist]

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 resolves a critical token-dimension mismatch occurring when using Qwen3.5 MTP with expert parallelism and upstream MoE sequence parallelism. By removing an outdated communication bypass that incorrectly performed a TP AllReduce instead of an EP ReduceScatter, the implementation now ensures consistent token layout across all parallel configurations. The changes improve reliability for speculative decoding in complex distributed setups without introducing new configuration flags.

Highlights

• Removal of Legacy Bypass: Removed the obsolete draft-VL TP AllReduce bypass from the EP sequence-parallel finalize path to prevent token-layout mismatches.
• Regression Coverage: Added a new unit test to ensure that draft-VL execution correctly adheres to the EP ReduceScatter contract.
• E2E Test Update: Updated the Qwen3.5 MTP four-card test case to explicitly cover DP, TP, EP, and upstream MoE sequence parallelism configurations.

New Features

🧠 You can now enable Memory (public preview) 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 the 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 counterproductive. 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.

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

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

---

Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [BugFix] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

• Fixes #<issue_number>
• Closes #<issue_number>
• Resolves #<issue_number>
• Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

[gemini-code-assist]

Code Review

Suggested PR Title:

[Ops][Misc] Simplify EP sequence-parallel finalize and update Qwen3.5 test

Suggested PR Summary:

### What this PR does / why we need it?
This PR simplifies the EP sequence-parallel finalize implementation in `vllm_ascend/ops/register_custom_ops.py` by removing the special-case handling for draft VL models that used `tensor_model_parallel_all_reduce`. It also updates the Qwen3.5 E2E test configuration to use DP2/TP2/EP/SP and adds a unit test to verify the EP layout is kept.

Feedback is provided regarding the use of `enable_flashcomm1` as a temporary compatibility marker in the E2E test, which could be brittle and affect maintainability.

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

### How was this patch tested?
Tested with new unit test `test_sp_ep_reduce_scatter_draft_vl_keeps_ep_layout` and updated E2E tests.

[github-actions]

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

[github-actions]

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

[nwpu-zxr]

Approve.