[Bugfix] Fix for 24530. Fix naive all2all shared expert overlap.

[bnellnm]

Purpose

Fix shared expert overlap with naive all2all backend. Fixes #24530. We were running the shared experts with the dispatched hidden states instead of the original hidden states.

Test Plan

Ran deepseek with VLLM_ALL2ALL_BACKEND=naive

Test Result

works

cc @yewentao256 , @tlrmchlsmth , @simon-mo

[gemini-code-assist]

Code Review

This pull request addresses a bug concerning shared expert overlap when using the naive all-to-all backend. The core of the fix involves reordering the dispatch operation to execute after the shared expert computation, which correctly ensures that shared experts operate on the original hidden_states. Additionally, a do_combine flag has been introduced to the reduce_output function, allowing the combine operation to be skipped for shared expert outputs, which is appropriate since they are not part of the dispatch-combine communication flow. The changes are logical, well-implemented, and include a new assertion for improved robustness.

[yewentao256]

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.96|±  |0.0197|
|     |       |strict-match    |     5|exact_match|↑  | 0.96|±  |0.0197|

Verified that this fixes the issue, thanks for the work!