feat: allow shared expert overlapping for FlashInfer one-sided all-to-all (#50569)
Signed-off-by: jdebache <jdebache@nvidia.com>
Julien Debache committed
Aug 11, 2026 at 19:06 UTC
52be12cfac0c5a18ba906814b2d2bcadb40a9c4b
1 file changed
+9
-2
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+9
-2
@@ -78,15 +78,22 @@ class SharedExperts(torch.nn.Module):
78
@property
79
def _disable_shared_experts_overlap(self) -> bool:
80
# Disable shared expert overlap if:
81
- # - we are using eplb with non-default backend, because of correctness issues
81
+ # - we are using eplb with non-safe backend, because of correctness issues
82
# - we are using flashinfer with DP, since there nothing to gain
83
+
84
+ # Both these comm backends have been shown to be safe for shared expert overlap.
85
+ _EPLB_OVERLAP_SAFE_BACKENDS = (
86
+ "allgather_reducescatter",
87
+ "flashinfer_nvlink_one_sided",
88
+ )
89
+
90
parallel_config = self._moe_config.moe_parallel_config
91
if getattr(self._layer, "shard_sequence_parallel", False):
92
# TODO: we may enable this to optimize further
93
return True
94
return (
95
parallel_config.enable_eplb
89
- and parallel_config.all2all_backend != "allgather_reducescatter"
96
+ and parallel_config.all2all_backend not in _EPLB_OVERLAP_SAFE_BACKENDS
97
) or parallel_config.use_fi_nvl_two_sided_kernels
98
99
def _determine_shared_experts_order(