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, 16:15 UTC
fused_moe: treat NIXL EP as batched experts (#40412)
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Itay Alroy committed
Apr 24, 2026 at 16:05 UTC
c9d3c6e6af7fb848d3f03e256484f68a00201020
4 files changed
+12
-9
vllm/model_executor/layers/fused_moe/config.py
+8
@@ -993,6 +993,10 @@ class FusedMoEParallelConfig:
993
def use_batched_activation_format(self):
994
return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels
995
996
+ @property
997
+ def needs_round_robin_routing_tables(self):
998
+ return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels
999
+
1000
@property
1001
def use_ag_rs_all2all_kernels(self):
1002
return (
@@ -1294,3 +1298,7 @@ class FusedMoEConfig:
1298
@property
1299
def use_nixl_ep_kernels(self):
1300
return self.moe_parallel_config.use_nixl_ep_kernels
1301
+
1302
+ @property
1303
+ def needs_round_robin_routing_tables(self):
1304
+ return self.moe_parallel_config.needs_round_robin_routing_tables
vllm/model_executor/layers/fused_moe/layer.py
+2
-4
@@ -180,8 +180,7 @@ def determine_expert_placement_strategy(
180
return "linear"
181
if (
182
moe_parallel_config.use_all2all_kernels
183
- and not moe_parallel_config.use_deepep_ll_kernels
184
- and not moe_parallel_config.use_nixl_ep_kernels
183
+ and not moe_parallel_config.needs_round_robin_routing_tables
184
):
185
logger.warning(
186
"Round-robin expert placement currently only supports "
@@ -687,8 +686,7 @@ class FusedMoE(PluggableLayer):
686
# Currently routing_tables only needed for round-robin expert placement
687
# with DeepEP-ll or NIXL EP all2all backends.
688
if self.expert_placement_strategy != "round_robin" or (
690
- not self.moe_parallel_config.use_deepep_ll_kernels
691
- and not self.moe_parallel_config.use_nixl_ep_kernels
689
+ not self.moe_parallel_config.needs_round_robin_routing_tables
690
):
691
return None
692
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+1
-1
@@ -884,7 +884,7 @@ def make_mxfp4_moe_kernel(
884
experts,
885
shared_experts=(
886
shared_experts
887
- if moe_config.moe_parallel_config.use_deepep_ll_kernels
887
+ if moe_config.moe_parallel_config.use_batched_activation_format
888
else None
889
),
890
inplace=(
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+1
-4
@@ -168,10 +168,7 @@ def select_nvfp4_moe_backend(
168
NvFp4MoeBackend.EMULATION,
169
]
170
171
- # NOTE(rob): this is kind of a hack. We need to peak into
172
- # the prepare-finalize selection to determine if we are using
173
- # the batched or standard expert format.
174
- use_batched = config.moe_parallel_config.use_deepep_ll_kernels
171
+ use_batched = config.moe_parallel_config.use_batched_activation_format
172
activation_format = (
173
mk.FusedMoEActivationFormat.BatchedExperts
174
if use_batched