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 14, 05:30 UTC
[Bugfix] Fix for 24530. Fix naive all2all shared expert overlap. (#24538)
bnellnm committed
Sep 9, 2025 at 20:53 UTC
b23fb78623f0019854af7a820d398cba9a316677
1 file changed
+9
-6
vllm/model_executor/layers/fused_moe/layer.py
+9
-6
@@ -1755,9 +1755,6 @@ class FusedMoE(CustomOp):
1755
self.dp_size > 1
1756
and not self.moe_parallel_config.use_deepep_ht_kernels
1757
and not self.moe_config.use_flashinfer_cutlass_kernels)
1758
- if do_naive_dispatch_combine:
1759
- hidden_states, router_logits = get_ep_group().dispatch(
1760
- hidden_states, router_logits)
1758
1759
# If there are shared experts but we are not using a modular kernel, the
1760
# shared experts must be called here
@@ -1768,6 +1765,10 @@ class FusedMoE(CustomOp):
1765
else:
1766
shared_output = None
1767
1768
+ if do_naive_dispatch_combine:
1769
+ hidden_states, router_logits = get_ep_group().dispatch(
1770
+ hidden_states, router_logits)
1771
+
1772
# Matrix multiply.
1773
final_hidden_states = self.quant_method.apply(
1774
layer=self,
@@ -1800,8 +1801,9 @@ class FusedMoE(CustomOp):
1801
final_hidden_states,
1802
)
1803
1803
- def reduce_output(states: torch.Tensor) -> torch.Tensor:
1804
- if do_naive_dispatch_combine:
1804
+ def reduce_output(states: torch.Tensor,
1805
+ do_combine: bool = True) -> torch.Tensor:
1806
+ if do_naive_dispatch_combine and do_combine:
1807
states = get_ep_group().combine(states)
1808
1809
if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
@@ -1810,10 +1812,11 @@ class FusedMoE(CustomOp):
1812
return states
1813
1814
if self.shared_experts is None:
1815
+ assert not isinstance(final_hidden_states, tuple)
1816
return reduce_output(final_hidden_states)
1817
else:
1818
return (
1816
- reduce_output(final_hidden_states[0]),
1819
+ reduce_output(final_hidden_states[0], do_combine=False),
1820
reduce_output(final_hidden_states[1]),
1821
)
1822