[BugFix][MoE] Keep shared and final maybe_all_reduce decisions runtime-aware

[drslark]

What this PR does / why we need it?

This PR reverts #16495 and replaces the complete MoE forward wrapper with runtime-aware reduction at three points: routed output before a transform or shared DP addition, shared TP output before addition, and final output after addition.

flowchart TD
    R["Routed output"] --> R1["1. Maybe reduce routed early<br/>Required by transform or shared DP?<br/>If yes, reduce unless already reduced"]
    R1 --> T["Routed transform, if any"]
    S["Shared output"] --> S1["2. Maybe reduce shared<br/>Shared TP + routed already reduced?<br/>If yes, reduce shared"]
    T --> A["Add shared and routed, if shared exists"]
    S1 --> A
    A --> F["3. Maybe reduce final output<br/>No SP + no early reduction<br/>+ routed not already reduced?<br/>If yes, reduce combined output"]
    F --> O["Return"]

Loading

Reduction conditions

Step	When to check	When to all-reduce
1. Routed early	Shared DP only, or a routed transform exists with SP disabled and (TP > 1 or EP > 1)	Routed output is not already reduced
2. Shared	Shared output exists and uses TENSOR_PARALLEL	Routed output was reduced early or already reduced by communication
3. Final	SP is disabled and early routed reduction is not required	Routed output is not already reduced; reduce after adding shared and routed outputs

Runtime check

Communication path	Routed already reduced?
MC2 / ALLTOALL / FUSED_MC2	Yes
ALLGATHER + SP	Yes
ALLGATHER without SP	No

Static vs. runtime decisions

Location	Decision	Evaluated at
The three _maybe_reduce_* runner methods	Layer configuration: shared layout, SP, transform presence, TP / EP sizes	FX tracing; fixed for the layer
Custom op real implementations	Whether routed output is already reduced, using the current communication type	Runtime
fake_impl	Output metadata only; no communication decision or collective	FX tracing

The runner methods select reduction points using static conditions. Communication-dependent decisions stay inside the custom ops and are evaluated by their real implementations at runtime. Both ops receive layer_name and read the layer's MoE configuration internally.

Does this PR introduce any user-facing change?

Fixes incorrect MoE output when tracing and replay use different communication types. No configuration changes are required.

How was this patch tested?

Validated the runtime reduction patch from eaca0c24d on the native 20260914 Ascend image. Only the reduction changes in fused_moe.py and register_custom_ops.py were applied; the rest of the image was retained. This does not validate the entire PR branch against current main.

Test	Result
AIS Bench GSM8K, first 20 test samples, concurrency 1	20/20 correct, 100.00% accuracy
AIS Bench GSM8K, the same 20 samples, concurrency 4	18/20 correct, 90.00% accuracy
Review of all 20 complete responses at concurrency 4	No looping repetition, garbled text or visibly truncated answers observed
Three original arithmetic reproducer prompts	3/3 correct: 391, 16, 150
Graph-replay regression cases extracted from the PR	10 passed, covering ALLGATHER, MC2, ALLTOALL and FUSED_MC2

The two incorrect responses at concurrency 4 contain complete reasoning and final answers:

Sample (1-based)	Predicted	Reference	Observed reasoning error
3	25000	70000	Applies the house-value increase percentage to the repair cost
13	12	13	Chooses the break-even year instead of the first profitable year

The concurrency-dependent accuracy difference has not been isolated to a root cause. The absence of looping or garbled output does not establish batch-invariant accuracy.

Graph mode was enabled: FULL_DECODE_ONLY.

Only AIS Bench concurrency changed between the two runs; the dataset, generation settings and running service were retained. The full PR test suite and full lint workflow have not been rerun on the latest revision.

• 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 addresses a critical issue in Mixture-of-Experts (MoE) model inference where graph tracing and replay could result in mismatched communication types, leading to incorrect output. By decoupling the reduction logic from static tracing decisions and introducing a runtime-aware custom operation, the system now correctly manages shared-expert and final tensor-parallel reductions based on the live communication context.

Highlights

• MoE Reduction Logic: Implemented a new custom operation maybe_all_reduce_tensor_model_parallel to handle shared and final reductions, ensuring they remain runtime-aware during graph replay.
• Communication Consistency: Fixed incorrect MoE output by ensuring that shared-expert reductions and final TP reductions use complementary runtime conditions, preventing corruption when tracing and replay communication types differ.
• Test Coverage: Added comprehensive graph-replay regression tests covering various communication types (ALLGATHER, MC2, ALLTOALL, FUSED_MC2) and MoE layouts to validate the fix.

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][Feature] Register maybe_all_reduce_tensor_model_parallel custom op for MoE reduction tracing

Suggested PR Summary:

### What this PR does / why we need it?
This PR registers a new custom operator `maybe_all_reduce_tensor_model_parallel` to handle conditional tensor model parallel all-reductions dynamically during graph replay, avoiding static branching on tracing-specific variables. It updates the Ascend MoE runner and fused MoE module to utilize this operator, ensuring correct reduction behavior across different communication configurations.

Additionally, a review feedback points out that the `fake_impl` for the registered custom op should return `states.clone()` instead of `states` to prevent PyTorch's `torch.compile` tracing from making incorrect aliasing assumptions.

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

### How was this patch tested?
The changes are tested with new and updated unit tests in `tests/ut/ops/test_fused_moe.py`, including a new test `test_maybe_all_reduce_graph_replay_preserves_shared_and_routed_outputs` that verifies graph replay behavior.

[gemini-code-assist]

Returning the input tensor states directly in the fake_impl of maybe_all_reduce_tensor_model_parallel causes PyTorch's functionalization and aliasing analysis during torch.compile tracing to assume that the output always aliases the input.

However, at runtime, _maybe_all_reduce_tensor_model_parallel_impl can perform an out-of-place tensor_model_parallel_all_reduce(states) which returns a new tensor with different storage. This mismatch between compile-time aliasing assumptions and runtime behavior can lead to subtle correctness bugs, compilation failures, or incorrect buffer reuse optimizations under torch.compile.

To prevent this, the fake_impl should return a new tensor that does not alias the input, such as states.clone().

Suggested change

	direct_register_custom_op(
	op_name="maybe_all_reduce_tensor_model_parallel",
	op_func=_maybe_all_reduce_tensor_model_parallel_impl,
	fake_impl=lambda states, is_sequence_parallel=False, reduce_when_reduced=False: states,
	mutates_args=[],
	dispatch_key="PrivateUse1",
	)
	direct_register_custom_op(
	op_name="maybe_all_reduce_tensor_model_parallel",
	op_func=_maybe_all_reduce_tensor_model_parallel_impl,
	fake_impl=lambda states, is_sequence_parallel=False, reduce_when_reduced=False: states.clone(),
	mutates_args=[],
	dispatch_key="PrivateUse1",
	)

[github-actions]

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

[drslark]

/rerun
[Bot]: rerun completed.

Rerun (failed jobs only):

• E2E