A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 17:00 UTC
[Feature] Enable TP + EP `shared_experts` overlap with router, 3.7% E2E performance improvement (#28164)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Nov 5, 2025 at 20:21 UTC
d71af5f5020e0fee5375b3cf7898852abbae22f2
2 files changed
+16
-8
vllm/model_executor/layers/fused_moe/layer.py
+1
-1
@@ -1178,7 +1178,7 @@ class FusedMoE(CustomOp):
1178
hidden_size: Input hidden state size of the transformer
1179
intermediate_size: Intermediate size of the experts
1180
params_dtype: Data type for the parameters.
1181
- reduce_results: Whether to all all_reduce on the output of the layer
1181
+ reduce_results: Whether to all_reduce on the output of the layer
1182
renormalize: Whether to renormalize the logits in the fused_moe kernel
1183
quant_config: Quantization configure.
1184
enable_eplb: Whether to enable expert parallelism load balancer.
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
+15
-7
@@ -3,7 +3,10 @@
3
4
import torch
5
6
-from vllm.distributed import tensor_model_parallel_all_reduce
6
+from vllm.distributed import (
7
+ get_tensor_model_parallel_world_size,
8
+ tensor_model_parallel_all_reduce,
9
+)
10
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
11
12
@@ -25,16 +28,13 @@ class SharedFusedMoE(FusedMoE):
28
super().__init__(**kwargs)
29
self._shared_experts = shared_experts
30
28
- # Disable shared expert overlap if EP is disabled or we are not using
31
+ # Disable shared expert overlap if we are not using
32
# flashinfer + DP since there is nothing to be gained in this case.
33
# Disabling the overlap optimization also prevents the shared experts
34
# from being hidden from torch.compile.
35
self.use_overlapped = (
36
use_overlapped
34
- and not (
35
- self.use_ep
36
- or (self.use_flashinfer_cutlass_kernels and self.dp_size > 1)
37
- )
37
+ and not (self.use_flashinfer_cutlass_kernels and self.dp_size > 1)
38
and self._shared_experts is not None
39
)
40
@@ -65,7 +65,7 @@ class SharedFusedMoE(FusedMoE):
65
# should have been created with reduce_results=False.
66
if (
67
self.reduce_results
68
- and self.tp_size > 1
68
+ and get_tensor_model_parallel_world_size() > 1
69
and self.must_reduce_shared_expert_outputs()
70
):
71
shared_out = tensor_model_parallel_all_reduce(shared_out)
@@ -81,4 +81,12 @@ class SharedFusedMoE(FusedMoE):
81
hidden_states=hidden_states,
82
router_logits=router_logits,
83
)
84
+ # ensure early TP reduction of shared expert outputs when required
85
+ if (
86
+ shared_out is not None
87
+ and self.reduce_results
88
+ and get_tensor_model_parallel_world_size() > 1
89
+ and self.must_reduce_shared_expert_outputs()
90
+ ):
91
+ shared_out = tensor_model_parallel_all_reduce(shared_out)
92
return shared_out, fused_out