Fix Qwen3 MoE: also guard EP all-reduce with not use_reduce_scatter (follow-up to #23731)

[ByronHsu]

Motivation

Follow-up to #23731 (which fixed double-reduce when DP attention + EP + reduce_scatterv).

The TP branch in Qwen3MoeSparseMoeBlock.forward_normal already guards on not use_reduce_scatter — the older LayerCommunicator-driven reduce-scatter path, which is distinct from the should_use_dp_reduce_scatterv() path added by #22642. The EP branch was missing this guard, so when LayerCommunicator's post-attention scatter does reduce-scatter, the EP all-reduce double-reduces and corrupts logits.

This is the same root cause as #23729 / #23731, just on the older use_reduce_scatter codepath instead of the newer should_use_dp_reduce_scatterv one. Both guards belong on the EP branch (and the TP branch already has both).

Modifications

Add not use_reduce_scatter to the EP all-reduce guard, mirroring the TP branch.

if (
    self.ep_size > 1
    and not should_allreduce_fusion
    and not use_reduce_scatter   # added
    and not should_use_dp_reduce_scatterv()
):
    final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states)

Checklist

• Format your code according to the Code Formatting with Pre-Commit
• Add unit tests as outlined in the Running Unit Tests
• Update documentation / docstrings / example tutorials as needed

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request updates the conditional logic in the forward_normal method of the Qwen3 MoE model to include use_reduce_scatter when determining whether to perform an expert parallel all-reduce. The review feedback suggests further extending this guard condition to include a check for should_use_flashinfer_cutlass_moe_fp4_allgather, which would align the expert parallel logic with the tensor parallel branch and prevent potential double-reduction issues when specialized communication paths are active.

[gemini-code-assist]

To fully align the EP all-reduce guard with the TP all-reduce guard below (lines 344-349), you should also include not should_use_flashinfer_cutlass_moe_fp4_allgather(). This ensures that if the FlashInfer specialized communication path is active, the EP all-reduce is correctly skipped to avoid double-reduction, maintaining consistency between the TP and EP branches.

            and not should_allreduce_fusion
            and not use_reduce_scatter
            and not should_use_flashinfer_cutlass_moe_fp4_allgather()
            and not should_use_dp_reduce_scatterv()