[Perf] Fix moe reduce_scatter perf regression by removing additional comm, 5% E2E throughput gain back.

[yewentao256]

Purpose

Part of #46654

Following up PR for #48036 and alternative for #48657

Originally:

# before
hidden [T/P,H] + residual [T/P,H]
        ↓ all gather
      [T/P,2H]
        ↓ all-gather
       [T,2H]
        ↓ split
hidden [T,H] + residual [T,H]
        ↓ all-reduce hidden
        ↓ final RMSNorm

# Now
hidden [T/P,H]
        ↓ * P (as TP hidden states would be duplicated)
        ↓  final RMSNorm with local residual
       [T/P,H]
        ↓ all-gather
        [T,H]

So all-gather(2H) + all-reduce(H) -> all-gather(H)

Test

vllm serve zai-org/GLM-5.2-FP8 --kv-cache-dtype fp8_e4m3 --enable-expert-parallel --tensor-parallel-size 4 --tool-call-parser glm47 --enable-auto-tool-choice --reasoning-parser glm45 --port 9256 --profiler-config.profiler=torch --profiler-config.torch_profiler_dir=/home/yewentao256/profile_vllm --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

vllm bench serve --model zai-org/GLM-5.2-FP8  --dataset-name random --host 127.0.0.1 --port 9256 --random-input-len 8192 --random-output-len 1 --request-rate inf --num-prompts 128 --num-warmups 16

# Now

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  3.08      
Total input tokens:                      1048576   
Total generated tokens:                  128       
Request throughput (req/s):              41.50     
Output token throughput (tok/s):         41.50     
Peak output token throughput (tok/s):    76.00     
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          340041.81 
---------------Time to First Token----------------
Mean TTFT (ms):                          1669.15   
Median TTFT (ms):                        1669.02   
P99 TTFT (ms):                           2724.38   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          0.00      
Median TPOT (ms):                        0.00      
P99 TPOT (ms):                           0.00      
---------------Inter-token Latency----------------
Mean ITL (ms):                           0.00      
Median ITL (ms):                         0.00      
P99 ITL (ms):                            0.00      
==================================================
============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  2.75      
Total input tokens:                      1048576   
Total generated tokens:                  128       
Request throughput (req/s):              46.61     
Output token throughput (tok/s):         46.61     
Peak output token throughput (tok/s):    75.00     
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          381839.84 
---------------Time to First Token----------------
Mean TTFT (ms):                          1452.11   
Median TTFT (ms):                        1315.32   
P99 TTFT (ms):                           2478.83   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          0.00      
Median TPOT (ms):                        0.00      
P99 TPOT (ms):                           0.00      
---------------Inter-token Latency----------------
Mean ITL (ms):                           0.00      
Median ITL (ms):                         0.00      
P99 ITL (ms):                            0.00      
==================================================

# main
============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  3.23      
Total input tokens:                      1048576   
Total generated tokens:                  128       
Request throughput (req/s):              39.63     
Output token throughput (tok/s):         39.63     
Peak output token throughput (tok/s):    50.00     
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          324671.20 
---------------Time to First Token----------------
Mean TTFT (ms):                          1668.69   
Median TTFT (ms):                        1649.00   
P99 TTFT (ms):                           3126.43   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          0.00      
Median TPOT (ms):                        0.00      
P99 TPOT (ms):                           0.00      
---------------Inter-token Latency----------------
Mean ITL (ms):                           0.00      
Median ITL (ms):                         0.00      
P99 ITL (ms):                            0.00      
==================================================
============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Benchmark duration (s):                  3.03      
Total input tokens:                      1048576   
Total generated tokens:                  128       
Request throughput (req/s):              42.26     
Output token throughput (tok/s):         42.26     
Peak output token throughput (tok/s):    60.00     
Peak concurrent requests:                128.00    
Total token throughput (tok/s):          346261.13 
---------------Time to First Token----------------
Mean TTFT (ms):                          1483.22   
Median TTFT (ms):                        1461.62   
P99 TTFT (ms):                           2754.37   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          0.00      
Median TPOT (ms):                        0.00      
P99 TPOT (ms):                           0.00      
---------------Inter-token Latency----------------
Mean ITL (ms):                           0.00      
Median ITL (ms):                         0.00      
P99 ITL (ms):                            0.00      
==================================================

lm_eval --model local-completions --model_args "base_url=http://127.0.0.1:9256/v1/completions,model=$MODEL,num_concurrent=1024" --tasks gsm8k

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

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yewentao256.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tlrmchlsmth]

I don't see why we would do hidden_states = hidden_states * get_tensor_model_parallel_world_size(). Perhaps the scalar coefficient is being taken care of by a subsequent norm?

Suggested change

	if is_sequence_parallel:
	hidden_states = hidden_states * get_tensor_model_parallel_world_size()
	else:
	hidden_states = tensor_model_parallel_all_reduce(hidden_states)
	if not is_sequence_parallel:
	hidden_states = tensor_model_parallel_all_reduce(hidden_states)

[yewentao256]

I took a further look. You are right, fixed, thanks!

[tlrmchlsmth]

nit: Why do we need both checks? I would think these are redundant

[yewentao256]

Nice catch, fixed

[tlrmchlsmth]

one nit, but otherwise LGTM now!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yewentao256.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yewentao256.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Documentation preview: https://vllm--48763.org.readthedocs.build/en/48763/