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 4, 22:30 UTC
[Bugfix] Replace code that disabled shared expert overlap (#39222)
Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
bnellnm committed
Apr 20, 2026 at 19:36 UTC
6867bcd076191e8a60f0e49e90059b01910725be
3 files changed
+17
-4
vllm/model_executor/layers/fused_moe/config.py
+1
@@ -990,6 +990,7 @@ class FusedMoEParallelConfig:
990
991
@property
992
def use_batched_activation_format(self):
993
+ # TODO(bnell): nixl also uses batched format
994
return self.use_deepep_ll_kernels
995
996
@property
vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py
+2
-4
@@ -427,9 +427,6 @@ class MoERunnerBase(MoERunner):
427
via the router, and the actual fused MoE computation. Returns
428
(shared_expert_output, fused_expert_output).
429
"""
430
- # Run this before quant_method to avoid inplace issues.
431
- # TODO(bnell): probably not needed anymore since inplace is
432
- # disabled when shared experts are present.
430
self._maybe_apply_shared_experts(
431
shared_experts_input, SharedExpertsOrder.NO_OVERLAP
432
)
@@ -447,7 +444,7 @@ class MoERunnerBase(MoERunner):
444
)
445
446
# Passing shared_experts_input in case SharedExpertsOrder is
450
- # NO_OVERLAP or MK_INTERNAL_OVERLAPPED.
447
+ # MK_INTERNAL_OVERLAPPED.
448
fused_out = self.quant_method.apply(
449
layer=layer,
450
x=hidden_states,
@@ -490,6 +487,7 @@ class MoERunnerBase(MoERunner):
487
# parallel execution of shared experts with the FusedMoE via
488
# separate cuda stream)
489
if self._shared_experts is not None:
490
+ assert shared_experts_input is not None
491
self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
492
493
def _maybe_add_zero_expert_output(
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+14
@@ -80,10 +80,24 @@ class SharedExperts:
80
"Enabled separate cuda stream for MoE shared_experts", scope="local"
81
)
82
83
+ @property
84
+ def _disable_shared_experts_overlap(self) -> bool:
85
+ # Disable shared expert overlap if:
86
+ # - we are using eplb with non-default backend, because of correctness issues
87
+ # - we are using flashinfer with DP, since there nothing to gain
88
+ parallel_config = self._moe_config.moe_parallel_config
89
+ return (
90
+ parallel_config.enable_eplb
91
+ and parallel_config.all2all_backend != "allgather_reducescatter"
92
+ ) or parallel_config.use_fi_nvl_two_sided_kernels
93
+
94
def _determine_shared_experts_order(
95
self,
96
hidden_states: torch.Tensor,
97
) -> SharedExpertsOrder:
98
+ if self._disable_shared_experts_overlap:
99
+ return SharedExpertsOrder.NO_OVERLAP
100
+
101
if self._quant_method.mk_owns_shared_expert:
102
return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
103