@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7cf56a59a2
Fork 0

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 5, 16:00 UTC

[MoE Refactor] Make SharedExperts class for use with DefaultMoERunner (#35153)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Apr 1, 2026 at 09:44 UTC 7cf56a59a267d2f32860c96a601e56bb29f6627a
34 files changed +556 -397
tests/kernels/moe/utils.py
-2
@@ -603,7 +603,6 @@ def make_shared_experts(
603 def modular_triton_fused_moe(
604 moe_config: FusedMoEConfig,
605 quant_config: FusedMoEQuantConfig,
606 - shared_experts: torch.nn.Module | None = None,
606 ) -> FusedMoEKernel:
607 return FusedMoEKernel(
608 maybe_make_prepare_finalize(
@@ -613,6 +612,5 @@ def modular_triton_fused_moe(
612 use_monolithic=False,
613 ),
614 TritonExperts(moe_config, quant_config),
616 - shared_experts,
615 inplace=False,
616 )
tools/ep_kernels/install_python_libraries.sh
+1
@@ -103,6 +103,7 @@ pushd "$WORKSPACE"
103 echo "Downloading NVSHMEM ${NVSHMEM_VER} for ${NVSHMEM_SUBDIR} ..."
104 curl -fSL "${NVSHMEM_URL}" -o "${NVSHMEM_FILE}"
105 tar -xf "${NVSHMEM_FILE}"
106 +rm -rf nvshmem
107 mv "${NVSHMEM_FILE%.tar.xz}" nvshmem
108 rm -f "${NVSHMEM_FILE}"
109 rm -rf nvshmem/lib/bin nvshmem/lib/share
vllm/distributed/elastic_ep/elastic_execute.py
+1 -2
@@ -410,8 +410,7 @@ class ElasticEPScalingExecutor:
410 # for the new EP size by resetting quant_method to base
411 for module in moe_modules:
412 if hasattr(module.quant_method, "old_quant_method"):
413 - module.quant_method = module.quant_method.old_quant_method
414 - module.runner = module._init_runner()
413 + module._replace_quant_method(module.quant_method.old_quant_method)
414 prepare_communication_buffer_for_model(self.worker.model_runner.model)
415
416 eplb_model_state.communicator = create_eplb_communicator(
vllm/lora/layers/fused_moe.py
-4
@@ -595,10 +595,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
595 def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs):
596 return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs)
597
598 - @property
599 - def _shared_experts(self):
600 - return self.base_layer._shared_experts
601 -
598 @property
599 def quant_method(self):
600 return self.base_layer.quant_method
vllm/model_executor/layers/fused_moe/config.py
+9
@@ -937,6 +937,15 @@ class FusedMoEParallelConfig:
937 all2all_backend: str # all2all backend for MoE communication
938 enable_eplb: bool # whether to enable expert load balancing
939
940 + @property
941 + def use_dp_chunking(self) -> bool:
942 + return (
943 + self.use_deepep_ll_kernels
944 + or self.use_mori_kernels
945 + or self.use_fi_nvl_two_sided_kernels
946 + or self.use_nixl_ep_kernels
947 + ) and envs.VLLM_ENABLE_MOE_DP_CHUNK
948 +
949 @property
950 def is_sequence_parallel(self) -> bool:
951 return self.sp_size > 1
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+2
@@ -1194,6 +1194,8 @@ def cutlass_moe_w4a8_fp8(
1194 quant_config=quant_config,
1195 group_size=group_size,
1196 ),
1197 + shared_experts=None,
1198 + inplace=False,
1199 )
1200
1201 return fn.apply(
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+1
@@ -53,6 +53,7 @@ class TrtLlmFp8ExpertsBase:
53 self.local_num_experts = moe_config.num_local_experts
54 self.ep_rank = moe_config.moe_parallel_config.ep_rank
55
56 + self.moe_config = moe_config
57 self.quant_config = quant_config
58
59 @staticmethod
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+3 -5
@@ -40,9 +40,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
40 def mk_owns_shared_expert(self) -> bool:
41 # NOTE(rob): temporary attribute to indicate support for
42 # completed migration to the new internal MK interface.
43 - return (
44 - self.moe_kernel is not None and self.moe_kernel.shared_experts is not None
45 - )
43 + return self.moe_kernel is not None and self.moe_kernel.owns_shared_experts
44
45 @abstractmethod
46 def create_weights(
@@ -163,7 +161,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
161 topk_weights: torch.Tensor,
162 topk_ids: torch.Tensor,
163 shared_experts_input: torch.Tensor | None,
166 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
164 + ) -> torch.Tensor:
165 raise NotImplementedError
166
167 def apply_monolithic(
@@ -171,5 +169,5 @@ class FusedMoEMethodBase(QuantizeMethodBase):
169 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
170 x: torch.Tensor,
171 router_logits: torch.Tensor,
174 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
172 + ) -> torch.Tensor:
173 raise NotImplementedError
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+6 -4
@@ -16,6 +16,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
16 FusedMoEKernel,
17 FusedMoEPrepareAndFinalizeModular,
18 )
19 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
20 + SharedExperts,
21 +)
22
23 logger = init_logger(__name__)
24
@@ -44,7 +47,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
47 moe_layer: torch.nn.Module,
48 old_quant_method: FusedMoEMethodBase,
49 prepare_finalize: FusedMoEPrepareAndFinalizeModular,
47 - shared_experts: torch.nn.Module | None,
50 + shared_experts: SharedExperts | None,
51 inplace: bool = False,
52 ) -> "FusedMoEModularMethod":
53 return FusedMoEModularMethod(
@@ -52,8 +55,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
55 FusedMoEKernel(
56 prepare_finalize,
57 old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
55 - shared_experts,
56 - moe_parallel_config=moe_layer.moe_parallel_config,
58 + shared_experts=shared_experts,
59 inplace=inplace,
60 ),
61 )
@@ -89,7 +91,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
91 topk_weights: torch.Tensor,
92 topk_ids: torch.Tensor,
93 shared_experts_input: torch.Tensor | None,
92 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
94 + ) -> torch.Tensor:
95 assert self.moe_kernel is not None
96 return self.moe_kernel.apply(
97 hidden_states=x,
vllm/model_executor/layers/fused_moe/layer.py
+23 -35
@@ -42,6 +42,9 @@ from vllm.model_executor.layers.fused_moe.router.router_factory import (
42 from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (
43 DefaultMoERunner,
44 )
45 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
46 + SharedExperts,
47 +)
48 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
49 UnquantizedFusedMoEMethod,
50 )
@@ -275,8 +278,6 @@ class FusedMoE(CustomOp):
278 ):
279 super().__init__()
280
278 - self._gate = gate
279 - self._shared_experts = shared_experts
281 self._routed_input_transform = routed_input_transform
282
283 if params_dtype is None:
@@ -486,7 +487,7 @@ class FusedMoE(CustomOp):
487 device=vllm_config.device_config.device,
488 routing_method=self.routing_method_type,
489 # TODO: in_dtype == out_dtype?
489 - disable_inplace=disable_inplace() or self._shared_experts is not None,
490 + disable_inplace=disable_inplace() or shared_experts is not None,
491 )
492 if self.moe_config.use_mori_kernels:
493 assert self.rocm_aiter_fmoe_enabled, (
@@ -564,34 +565,20 @@ class FusedMoE(CustomOp):
565 moe_quant_params["intermediate_size_full"] = intermediate_size
566
567 self.quant_method.create_weights(layer=self, **moe_quant_params)
567 - self.base_quant_method = self.quant_method
568 -
569 - # Disable shared expert overlap if:
570 - # - we are using eplb with non-default backend, because of correctness issues
571 - # - we are using flashinfer with DP, since there nothing to gain
572 - # - we are using marlin kernels
573 - backend = self.moe_parallel_config.all2all_backend
574 - self.use_overlapped = (
575 - not (
576 - (self.enable_eplb and backend != "allgather_reducescatter")
577 - or self.moe_parallel_config.use_fi_nvl_two_sided_kernels
578 - )
579 - and self._shared_experts is not None
580 - )
568
582 - self.runner = self._init_runner()
569 + # TODO(bnell): this is un-needed and removed in a follow up PR.
570 + self.base_quant_method = self.quant_method
571
584 - def _init_runner(self):
572 # Storing the runner in the FusedMoE is an intermediate state, eventually
573 # the runner will own the FusedMoE layer and provide the execution interface
574 # for MoE ops.
588 - return DefaultMoERunner(
575 + self.runner = DefaultMoERunner(
576 layer=self,
577 moe_config=self.moe_config,
578 router=self.router,
579 routed_input_transform=self._routed_input_transform,
593 - gate=self.gate,
594 - shared_experts=self.shared_experts,
580 + gate=gate,
581 + shared_experts=shared_experts,
582 quant_method=self.quant_method,
583 reduce_results=self.reduce_results,
584 enable_dbo=self.vllm_config.parallel_config.enable_dbo,
@@ -602,10 +589,7 @@ class FusedMoE(CustomOp):
589 # intrusive way to do this.
590 def _replace_quant_method(self, mk: FusedMoEMethodBase):
591 self.quant_method = mk
605 - # We need to force reconstruction of runner because we're swapping out
606 - # the quant_method with a FusedMoEModularMethod. This logic can go
607 - # away once the FusedMoEModularMethod is eliminated.
608 - self.runner = self._init_runner()
592 + self.runner._replace_quant_method(mk)
593
594 # Note: maybe_init_modular_kernel should only be called by
595 # prepare_communication_buffer_for_model.
@@ -639,8 +623,8 @@ class FusedMoE(CustomOp):
623 )
624
625 @property
642 - def shared_experts(self) -> torch.nn.Module | None:
643 - return self._shared_experts if self.use_overlapped else None
626 + def shared_experts(self) -> SharedExperts | None:
627 + return self.runner.shared_experts
628
629 @property
630 def layer_id(self):
@@ -649,10 +633,6 @@ class FusedMoE(CustomOp):
633
634 return extract_layer_index(self.layer_name)
635
652 - @property
653 - def gate(self) -> torch.nn.Module | None:
654 - return self._gate if self.use_overlapped else None
655 -
636 @property
637 def tp_size(self):
638 return self.moe_parallel_config.tp_size
@@ -676,7 +656,7 @@ class FusedMoE(CustomOp):
656 @property
657 def is_internal_router(self) -> bool:
658 # By default, router/gate is called before FusedMoE forward pass
679 - return self.gate is not None
659 + return self.runner.is_internal_router()
660
661 def _maybe_init_expert_routing_tables(
662 self,
@@ -1467,7 +1447,12 @@ class FusedMoE(CustomOp):
1447 assert all(
1448 weight.is_contiguous()
1449 for name, weight in weights
1470 - if not (name.startswith("_shared_experts.") or name.startswith("_gate."))
1450 + if not (
1451 + name.startswith("_shared_experts.")
1452 + or name.startswith("_gate.")
1453 + or name.startswith("_routed_input_transform.")
1454 + or name.startswith("_routed_output_transform.")
1455 + )
1456 and name not in NON_EXPERT_WEIGHTS
1457 )
1458
@@ -1477,8 +1462,11 @@ class FusedMoE(CustomOp):
1462 if name not in NON_EXPERT_WEIGHTS
1463 and weight.shape != torch.Size([])
1464 and not name.startswith("_shared_experts.")
1480 - # exclude parameters from non-expert submodules (e.g. gate/shared)
1465 + # exclude parameters from non-expert submodules,
1466 + # e.g. gate/shared/transforms.
1467 and not name.startswith("_gate.")
1468 + and not name.startswith("_routed_input_transform.")
1469 + and not name.startswith("_routed_output_transform.")
1470 ]
1471
1472 def set_eplb_state(
vllm/model_executor/layers/fused_moe/modular_kernel.py
+42 -35
@@ -21,6 +21,10 @@ from vllm.model_executor.layers.fused_moe.config import (
21 FusedMoEQuantConfig,
22 RoutingMethodType,
23 )
24 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
25 + SharedExperts,
26 + SharedExpertsOrder,
27 +)
28 from vllm.model_executor.layers.fused_moe.utils import (
29 _resize_cache,
30 disable_inplace,
@@ -235,6 +239,13 @@ class FusedMoEPrepareAndFinalize(ABC):
239 """
240 raise NotImplementedError
241
242 + def supports_async(self) -> bool:
243 + """
244 + Indicates whether or not this class implements prepare_async and
245 + finalize_async.
246 + """
247 + return False
248 +
249
250 # TODO: pass FusedMoEParallelConfig in as ctor parameter?
251 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
@@ -281,13 +292,6 @@ class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
292 """
293 raise NotImplementedError
294
284 - def supports_async(self) -> bool:
285 - """
286 - Indicates whether or not this class implements prepare_async and
287 - finalize_async.
288 - """
289 - return False
290 -
295 def prepare_async(
296 self,
297 a1: torch.Tensor,
@@ -1003,15 +1007,20 @@ class FusedMoEKernelModularImpl:
1007 self,
1008 prepare_finalize: FusedMoEPrepareAndFinalizeModular,
1009 fused_experts: FusedMoEExpertsModular,
1006 - shared_experts: torch.nn.Module | None = None,
1007 - moe_parallel_config: FusedMoEParallelConfig | None = None,
1010 + shared_experts: SharedExperts | None,
1011 inplace: bool = False,
1012 ):
1013 self.prepare_finalize = prepare_finalize
1014 self.fused_experts = fused_experts
1012 - self.shared_experts = shared_experts
1013 - self.moe_parallel_config = moe_parallel_config
1015 + # Only accept shared experts if they can be run w/async.
1016 + # The MoERunner/SharedExperts class will coordinate with the MK to ensure
1017 + # that the SharedExperts are executed only once.
1018 + self.shared_experts = (
1019 + shared_experts if prepare_finalize.supports_async() else None
1020 + )
1021 self.inplace = inplace
1022 + moe_parallel_config = fused_experts.moe_config.moe_parallel_config
1023 + self.moe_parallel_config = moe_parallel_config
1024 self.is_dp_ep = (
1025 moe_parallel_config is not None
1026 and moe_parallel_config.dp_size > 1
@@ -1081,6 +1090,17 @@ class FusedMoEKernelModularImpl:
1090
1091 return workspace13, workspace2, fused_out
1092
1093 + def _maybe_apply_shared_experts(
1094 + self,
1095 + shared_experts_input: torch.Tensor | None,
1096 + ):
1097 + if self.shared_experts is not None:
1098 + assert shared_experts_input is not None
1099 + self.shared_experts.apply(
1100 + shared_experts_input,
1101 + SharedExpertsOrder.MK_INTERNAL_OVERLAPPED,
1102 + )
1103 +
1104 def _prepare(
1105 self,
1106 hidden_states: torch.Tensor,
@@ -1253,15 +1273,6 @@ class FusedMoEKernelModularImpl:
1273 shared_experts_input is the original hidden_states (full
1274 dimension) needed by the shared expert MLP.
1275 """
1256 - shared_output: torch.Tensor | None = None
1257 -
1258 - # For latent MoE: shared experts need the original hidden_states
1259 - # (full hidden_size), not the latent-projected version used by
1260 - # routed experts.
1261 - se_hidden_states = (
1262 - shared_experts_input if shared_experts_input is not None else hidden_states
1263 - )
1264 -
1276 if not self.prepare_finalize.supports_async():
1277 assert not dbo_enabled()
1278
@@ -1273,8 +1284,6 @@ class FusedMoEKernelModularImpl:
1284 apply_router_weight_on_input,
1285 self.fused_experts.finalize_weight_and_reduce_impl(),
1286 )
1276 - if self.shared_experts is not None:
1277 - shared_output = self.shared_experts(se_hidden_states)
1287 else:
1288 finalize_ret = self.prepare_finalize.finalize_async(
1289 output,
@@ -1284,8 +1293,7 @@ class FusedMoEKernelModularImpl:
1293 apply_router_weight_on_input,
1294 self.fused_experts.finalize_weight_and_reduce_impl(),
1295 )
1287 - if self.shared_experts is not None:
1288 - shared_output = self.shared_experts(se_hidden_states)
1296 + self._maybe_apply_shared_experts(shared_experts_input)
1297
1298 # TODO(lucas): refactor this in the alternative schedules followup
1299 # currently unpack if we have hook + receiver pair or just
@@ -1308,11 +1316,7 @@ class FusedMoEKernelModularImpl:
1316
1317 receiver()
1318
1311 - if self.shared_experts is None:
1312 - return output
1313 - else:
1314 - assert shared_output is not None
1315 - return shared_output, output
1319 + return output
1320
1321 def apply(
1322 self,
@@ -1326,7 +1330,7 @@ class FusedMoEKernelModularImpl:
1330 expert_map: torch.Tensor | None = None,
1331 apply_router_weight_on_input: bool = False,
1332 shared_experts_input: torch.Tensor | None = None,
1329 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1333 + ) -> torch.Tensor:
1334 """
1335 This function computes a Mixture of Experts (MoE) layer using two sets
1336 of weights, w1 and w2, and top-k gating mechanism.
@@ -1469,12 +1473,10 @@ class FusedMoEKernel:
1473 self,
1474 prepare_finalize: FusedMoEPrepareAndFinalize,
1475 fused_experts: FusedMoEExperts,
1472 - shared_experts: torch.nn.Module | None = None,
1473 - moe_parallel_config: FusedMoEParallelConfig | None = None,
1476 + shared_experts: SharedExperts | None = None,
1477 inplace: bool = False,
1478 ):
1479 super().__init__()
1477 - self.shared_experts = shared_experts # NOTE: check if we can remove
1480
1481 # Initialize the implementation (monolithic or modular).
1482 self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl
@@ -1485,14 +1487,12 @@ class FusedMoEKernel:
1487 prepare_finalize,
1488 fused_experts,
1489 shared_experts,
1488 - moe_parallel_config,
1490 inplace,
1491 )
1492
1493 elif isinstance(
1494 prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic
1495 ) and isinstance(fused_experts, FusedMoEExpertsMonolithic):
1495 - assert shared_experts is None
1496 assert not inplace
1497 self.impl = FusedMoEKernelMonolithicImpl(
1498 prepare_finalize,
@@ -1508,6 +1508,13 @@ class FusedMoEKernel:
1508
1509 self._post_init_setup()
1510
1511 + @property
1512 + def owns_shared_experts(self) -> bool:
1513 + if isinstance(self.impl, FusedMoEKernelModularImpl):
1514 + return self.impl.shared_experts is not None
1515 + else:
1516 + return False
1517 +
1518 @property
1519 def is_monolithic(self) -> bool:
1520 return isinstance(self.impl, FusedMoEKernelMonolithicImpl)
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+5 -7
@@ -18,6 +18,9 @@ from vllm.model_executor.layers.fused_moe.config import (
18 fp8_w8a8_moe_quant_config,
19 fp8_w8a16_moe_quant_config,
20 )
21 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
22 + SharedExperts,
23 +)
24 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
25 FlashinferMoeBackend,
26 get_flashinfer_moe_backend,
@@ -545,7 +548,7 @@ def make_fp8_moe_kernel(
548 experts_cls: type[mk.FusedMoEExperts],
549 fp8_backend: Fp8MoeBackend,
550 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
548 - shared_experts: torch.nn.Module | None = None,
551 + shared_experts: SharedExperts | None = None,
552 ) -> mk.FusedMoEKernel:
553 # Create Prepare/Finalize.
554 prepare_finalize = maybe_make_prepare_finalize(
@@ -581,12 +584,7 @@ def make_fp8_moe_kernel(
584 kernel = mk.FusedMoEKernel(
585 prepare_finalize,
586 experts,
584 - shared_experts=(
585 - shared_experts
586 - if moe_config.moe_parallel_config.use_deepep_ll_kernels
587 - else None
588 - ),
589 - moe_parallel_config=moe_config.moe_parallel_config,
587 + shared_experts=shared_experts,
588 inplace=(
589 not moe_config.disable_inplace
590 and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
-1
@@ -859,7 +859,6 @@ def make_mxfp4_moe_kernel(
859 if moe_config.moe_parallel_config.use_deepep_ll_kernels
860 else None
861 ),
862 - moe_parallel_config=moe_config.moe_parallel_config,
862 inplace=(
863 not moe_config.disable_inplace and mxfp4_backend not in TRTLLM_BACKENDS
864 ),
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+5 -7
@@ -17,6 +17,9 @@ from vllm.model_executor.layers.fused_moe.config import (
17 nvfp4_moe_quant_config,
18 nvfp4_w4a16_moe_quant_config,
19 )
20 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
21 + SharedExperts,
22 +)
23 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
24 prepare_nvfp4_moe_layer_for_fi_or_cutlass,
25 )
@@ -386,7 +389,7 @@ def make_nvfp4_moe_kernel(
389 moe_config: FusedMoEConfig,
390 experts_cls: type[mk.FusedMoEExperts],
391 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
389 - shared_experts: torch.nn.Module | None = None,
392 + shared_experts: SharedExperts | None = None,
393 ) -> mk.FusedMoEKernel:
394 # Create Prepare/Finalize.
395 prepare_finalize = maybe_make_prepare_finalize(
@@ -422,12 +425,7 @@ def make_nvfp4_moe_kernel(
425 kernel = mk.FusedMoEKernel(
426 prepare_finalize,
427 experts,
425 - shared_experts=(
426 - shared_experts
427 - if moe_config.moe_parallel_config.use_deepep_ll_kernels
428 - else None
429 - ),
430 - moe_parallel_config=moe_config.moe_parallel_config,
428 + shared_experts=shared_experts,
429 inplace=False,
430 )
431
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+5 -7
@@ -18,6 +18,9 @@ from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEConfig,
19 FusedMoEQuantConfig,
20 )
21 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
22 + SharedExperts,
23 +)
24 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
25 FlashinferMoeBackend,
26 convert_moe_weights_to_flashinfer_trtllm_block_layout,
@@ -321,7 +324,7 @@ def make_unquantized_moe_kernel(
324 backend: UnquantizedMoeBackend,
325 experts_cls: type[mk.FusedMoEExperts],
326 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
324 - shared_experts: torch.nn.Module | None = None,
327 + shared_experts: SharedExperts | None = None,
328 ) -> mk.FusedMoEKernel:
329 # Create Prepare/Finalize
330 is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
@@ -355,12 +358,7 @@ def make_unquantized_moe_kernel(
358 kernel = mk.FusedMoEKernel(
359 prepare_finalize,
360 experts,
358 - shared_experts=(
359 - shared_experts
360 - if moe_config.moe_parallel_config.use_deepep_ll_kernels
361 - else None
362 - ),
363 - moe_parallel_config=moe_config.moe_parallel_config,
361 + shared_experts=shared_experts,
362 inplace=(not moe_config.disable_inplace and not is_monolithic),
363 )
364
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py
+1 -1
@@ -325,7 +325,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
325 **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
326 **(
327 dict(x_global_scale=qc_a1_gscale_or_scale)
328 - if qc_a1_gscale_or_scale is not None
328 + if qc_a1_gscale_or_scale is not None and nvfp4_dispatch
329 else dict()
330 ),
331 async_finish=False,
vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
+8
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 +from collections.abc import Callable
5
6 import torch
7
@@ -13,6 +14,13 @@ class FusedMoERouter(ABC):
14 method that is used for routing hidden states based on router logits.
15 """
16
17 + @abstractmethod
18 + def set_capture_fn(
19 + self,
20 + capture_fn: Callable[[torch.Tensor], None] | None,
21 + ) -> None:
22 + raise NotImplementedError
23 +
24 @property
25 @abstractmethod
26 def routing_method_type(self) -> RoutingMethodType:
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+183 -219
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING
7 import torch
8 import torch.nn.functional as F
9
10 -import vllm.envs as envs
10 from vllm.distributed import (
11 get_ep_group,
12 get_pcp_group,
@@ -29,13 +28,15 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
28 FusedMoERouter,
29 )
30 from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
31 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
32 + SharedExperts,
33 + SharedExpertsOrder,
34 +)
35 from vllm.platforms import current_platform
36 from vllm.utils.math_utils import cdiv
37 from vllm.utils.torch_utils import (
38 HAS_OPAQUE_TYPE,
39 ModuleName,
37 - aux_stream,
38 - current_stream,
40 direct_register_custom_op,
41 )
42 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
@@ -74,6 +75,9 @@ def _resolve_layer_name(layer_name: str | ModuleName) -> str:
75 return layer_name.value if isinstance(layer_name, ModuleName) else layer_name
76
77
78 +# Note: _moe_forward and _moe_forward_shared should not contain any
79 +# implementation details, They should merely pass along control to
80 +# the runner's 'forward_dispatch' method.
81 def _moe_forward(
82 hidden_states: torch.Tensor,
83 router_logits: torch.Tensor,
@@ -81,24 +85,12 @@ def _moe_forward(
85 layer_name: _layer_name_type,
86 ) -> torch.Tensor:
87 layer = get_layer_from_name(_resolve_layer_name(layer_name))
84 - # TODO(bnell): this can be removed after MK migration is complete.
85 - layer.ensure_moe_quant_config_init()
86 - runner = layer.runner
87 - with runner._sequence_parallel_context():
88 - if runner.use_dp_chunking:
89 - return runner.forward_impl_chunked(
90 - layer,
91 - hidden_states,
92 - router_logits,
93 - shared_experts_input,
94 - )
95 - else:
96 - return runner.forward_impl(
97 - layer,
98 - hidden_states,
99 - router_logits,
100 - shared_experts_input,
101 - )
88 + return layer.runner.forward_dispatch(
89 + layer,
90 + hidden_states,
91 + router_logits,
92 + shared_experts_input,
93 + )
94
95
96 def _moe_forward_fake(
@@ -117,24 +109,12 @@ def _moe_forward_shared(
109 layer_name: _layer_name_type,
110 ) -> tuple[torch.Tensor, torch.Tensor]:
111 layer = get_layer_from_name(_resolve_layer_name(layer_name))
120 - # TODO(bnell): this can be removed after MK migration is complete.
121 - layer.ensure_moe_quant_config_init()
122 - runner = layer.runner
123 - with runner._sequence_parallel_context():
124 - if runner.use_dp_chunking:
125 - return runner.forward_impl_chunked(
126 - layer,
127 - hidden_states,
128 - router_logits,
129 - shared_experts_input,
130 - )
131 - else:
132 - return runner.forward_impl(
133 - layer,
134 - hidden_states,
135 - router_logits,
136 - shared_experts_input,
137 - )
112 + return layer.runner.forward_dispatch(
113 + layer,
114 + hidden_states,
115 + router_logits,
116 + shared_experts_input,
117 + )
118
119
120 def _moe_forward_shared_fake(
@@ -159,7 +139,7 @@ def _moe_forward_shared_fake(
139 direct_register_custom_op(
140 op_name="moe_forward",
141 op_func=_moe_forward,
162 - mutates_args=["hidden_states"],
142 + mutates_args=["hidden_states"], # is this still true?
143 fake_impl=_moe_forward_fake,
144 tags=(torch.Tag.needs_fixed_stride_order,),
145 )
@@ -168,7 +148,6 @@ direct_register_custom_op(
148 direct_register_custom_op(
149 op_name="moe_forward_shared",
150 op_func=_moe_forward_shared,
171 - mutates_args=["hidden_states"],
151 fake_impl=_moe_forward_shared_fake,
152 tags=(torch.Tag.needs_fixed_stride_order,),
153 )
@@ -213,87 +192,68 @@ class DefaultMoERunner(MoERunner):
192 self.router = router
193 self.routed_input_transform = routed_input_transform
194 self.gate = gate
216 - self.shared_experts = shared_experts
195 self.quant_method = quant_method
196 self.reduce_results = reduce_results
197 self.enable_dbo = enable_dbo
198
199 + self.shared_experts: SharedExperts | None = None
200 + if shared_experts is not None:
201 + self.shared_experts = SharedExperts(
202 + shared_experts,
203 + moe_config=moe_config,
204 + # Note: For now we must pass quant_method along to SharedExperts so it
205 + # can property determine where the shared experts are supposed to be
206 + # called, i.e. by a MK or by the MoERunner.
207 + # Once the MK can be created upfront, we can just pass in the proper
208 + # flags derived from the quant_method's MK.
209 + reduce_results=reduce_results,
210 + quant_method=quant_method,
211 + enable_dbo=enable_dbo,
212 + )
213 +
214 # Chunked all2all staging tensor
222 - # TODO(bnell) rename these?
215 + # These need to exist ahead of time due to CUDAgraph construction
216 + # needing a fixed buffer address.
217 + self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking
218 self.batched_hidden_states: torch.Tensor | None = None
219 self.batched_router_logits: torch.Tensor | None = None
220 self._maybe_init_dp_chunking()
221
227 - # Allow disabling of the separate shared experts stream for
228 - # debug purposes.
229 - # TODO: Remove this after more extensive testings with TP/DP
230 - # and other execution modes
231 - self.use_shared_experts_stream = False
232 - if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:
233 - logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local")
234 - self.shared_experts_stream = None
235 - else:
236 - # TODO(rob): enable shared expert overlap with non-cuda-alike.
237 - # aux_stream() returns None on non-cuda-alike platforms.
238 - self.shared_experts_stream = aux_stream()
239 - if self.shared_experts_stream is not None:
240 - logger.debug_once(
241 - "Enabled separate cuda stream for MoE shared_experts", scope="local"
242 - )
243 -
222 # Needed for string -> FusedMoE layer lookup in custom ops.
223 self.layer_name = layer.layer_name
224
247 - self.moe_forward = self._select_forward(layer)
225 + self.forward_entry, self.forward_impl = self._select_forward(layer)
226 +
227 + def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]:
228 + # Select implementation based on presence of DP chunking.
229 + forward_impl_fn = (
230 + self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl
231 + )
232
249 - def _select_forward(self, layer: torch.nn.Module) -> Callable:
233 if current_platform.is_tpu() or current_platform.is_cpu():
234 # TODO: Once the OOM issue for the TPU backend is resolved, we
235 # will switch to using the moe_forward custom op.
236 # Note: CPU doesn't require wrapped forward_impl.
254 - return _moe_forward if self.shared_experts is None else _moe_forward_shared
237 + return (
238 + _moe_forward if self.shared_experts is None else _moe_forward_shared,
239 + forward_impl_fn,
240 + )
241
242 return (
243 torch.ops.vllm.moe_forward
244 if self.shared_experts is None
259 - else torch.ops.vllm.moe_forward_shared
245 + else torch.ops.vllm.moe_forward_shared,
246 + forward_impl_fn,
247 )
248
262 - @property
263 - def use_dp_chunking(self) -> bool:
264 - return (
265 - self.moe_config.moe_parallel_config.use_deepep_ll_kernels
266 - or self.moe_config.moe_parallel_config.use_mori_kernels
267 - or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels
268 - or self.moe_config.moe_parallel_config.use_nixl_ep_kernels
269 - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK
270 -
271 - def _maybe_setup_shared_experts_stream(
272 - self,
273 - hidden_states: torch.Tensor,
274 - shared_input: torch.Tensor | None,
275 - ):
276 - if self.use_shared_experts_stream:
277 - assert self.shared_experts_stream is not None
278 - assert self.moe_config.disable_inplace
279 -
280 - shared_experts_input = (
281 - shared_input if shared_input is not None else hidden_states
282 - )
283 -
284 - # Record that the shared_experts_input will be used in the
285 - # shared_experts_stream to avoid gc issue from
286 - # deallocation. For more details:
287 - # https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
288 - # NOTE: We don't need shared_output.record_stream(current_stream())
289 - # because we synch the streams before using shared_output.
290 - shared_experts_input.record_stream(self.shared_experts_stream)
249 + # TODO(bnell): temporary hack, do not call this method.
250 + def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
251 + if self.shared_experts is not None:
252 + self.shared_experts._quant_method = quant_method
253 + self.quant_method = quant_method
254
292 - # Mark sync start point for the separate shared experts
293 - # stream here since we want to run in parallel with the
294 - # router/gate (next op below)
295 - assert self.shared_experts_stream is not None
296 - self.shared_experts_stream.wait_stream(current_stream())
255 + def is_internal_router(self) -> bool:
256 + return self.gate is not None
257
258 def _maybe_init_dp_chunking(self):
259 if not self.use_dp_chunking:
@@ -325,38 +285,6 @@ class DefaultMoERunner(MoERunner):
285 device=device,
286 )
287
328 - @property
329 - def has_separate_shared_experts(self) -> bool:
330 - return (
331 - not self.quant_method.mk_owns_shared_expert
332 - and self.shared_experts is not None
333 - )
334 -
335 - def _apply_shared_experts(
336 - self,
337 - hidden_states: torch.Tensor,
338 - allow_streaming: bool = False,
339 - ) -> torch.Tensor | None:
340 - shared_output: torch.Tensor | None = None
341 - if self.has_separate_shared_experts:
342 - assert self.shared_experts is not None
343 -
344 - if self.use_shared_experts_stream and allow_streaming:
345 - # Run shared experts in parallel on a separate stream
346 - # NOTE: We start the separate stream here and mark the
347 - # sync end point immediately after it is done. This is
348 - # important to avoid excessive stream allocations by the cuda
349 - # graph replay later.
350 - with torch.cuda.stream(self.shared_experts_stream):
351 - # Note that hidden_states clone() is necessary here to avoid
352 - # conflict with the main stream
353 - shared_output = self.shared_experts(hidden_states)
354 - current_stream().wait_stream(self.shared_experts_stream)
355 - else:
356 - shared_output = self.shared_experts(hidden_states)
357 -
358 - return shared_output
359 -
288 def must_reduce_shared_expert_outputs(self) -> bool:
289 """
290 The shared_experts are typically computed using the RowParallelLinear
@@ -384,7 +312,9 @@ class DefaultMoERunner(MoERunner):
312 else:
313 return tensor_model_parallel_all_reduce(final_hidden_states)
314
387 - def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor:
315 + def apply_routed_input_transform(
316 + self, hidden_states: torch.Tensor
317 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
318 """Apply transform for routed experts (e.g., latent projection).
319
320 This is called by FusedMoE.forward_native. The original hidden_states
@@ -394,15 +324,22 @@ class DefaultMoERunner(MoERunner):
324 TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be
325 moved inside SharedFusedMoE to all-reduce on the smaller latent
326 dimension.
327 +
328 + Returns (possibly transformed) hidden states and the input for shared
329 + experts (or None if there are no shared experts).
330 """
331 if self.routed_input_transform is not None:
332 result = self.routed_input_transform(hidden_states)
333 # ReplicatedLinear returns (output, extra_bias) tuple.
334 # We only need the output tensor; extra_bias is not used here.
335 if isinstance(result, tuple):
403 - return result[0]
404 - return result
405 - return hidden_states
336 + return result[0], hidden_states
337 + return result, hidden_states
338 +
339 + return (
340 + hidden_states,
341 + hidden_states if self.shared_experts is not None else None,
342 + )
343
344 def _maybe_reduce_output(
345 self,
@@ -446,13 +383,11 @@ class DefaultMoERunner(MoERunner):
383
384 def _maybe_pad_hidden_states(
385 self,
449 - original_hidden_states: torch.Tensor | None,
386 + shared_experts_input: torch.Tensor | None,
387 hidden_states: torch.Tensor,
388 ) -> tuple[torch.Tensor, list[int]]:
452 - original_hidden_dim = (
453 - original_hidden_states.shape[-1]
454 - if original_hidden_states is not None
455 - else 0
389 + shared_experts_hidden_dim = (
390 + shared_experts_input.shape[-1] if shared_experts_input is not None else 0
391 )
392 transformed_hidden_dim = hidden_states.shape[-1]
393 if (
@@ -467,29 +402,37 @@ class DefaultMoERunner(MoERunner):
402 )
403
404 if self.shared_experts is not None:
470 - orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim]
405 + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim]
406 else:
407 orig_hidden_dims = [transformed_hidden_dim]
408
409 return hidden_states, orig_hidden_dims
410
411 + def _maybe_apply_shared_experts(
412 + self,
413 + shared_experts_input: torch.Tensor | None,
414 + order: SharedExpertsOrder,
415 + ):
416 + if self.shared_experts is not None:
417 + assert shared_experts_input is not None
418 + self.shared_experts.apply(shared_experts_input, order)
419 +
420 def _apply_quant_method(
421 self,
422 layer: torch.nn.Module,
423 hidden_states: torch.Tensor,
424 router_logits: torch.Tensor,
481 - shared_input: torch.Tensor | None,
482 - run_shared_experts_before: bool = True,
425 + shared_experts_input: torch.Tensor | None,
426 ) -> tuple[torch.Tensor | None, torch.Tensor]:
484 - shared_input = shared_input if shared_input is not None else hidden_states
485 - shared_output: torch.Tensor | None = None
486 -
427 # Run this before quant_method to avoid inplace issues.
488 - if run_shared_experts_before:
489 - shared_output = self._apply_shared_experts(shared_input, False)
428 + # TODO(bnell): probably not needed anymore since inplace is
429 + # disabled when shared experts are present.
430 + self._maybe_apply_shared_experts(
431 + shared_experts_input, SharedExpertsOrder.NO_OVERLAP
432 + )
433
434 if self.quant_method.is_monolithic:
492 - result = self.quant_method.apply_monolithic(
435 + fused_out = self.quant_method.apply_monolithic(
436 layer=layer,
437 x=hidden_states,
438 router_logits=router_logits,
@@ -500,25 +443,25 @@ class DefaultMoERunner(MoERunner):
443 router_logits=router_logits,
444 )
445
503 - result = self.quant_method.apply(
446 + # Passing shared_experts_input in case SharedExpertsOrder is
447 + # NO_OVERLAP or MK_INTERNAL_OVERLAPPED.
448 + fused_out = self.quant_method.apply(
449 layer=layer,
450 x=hidden_states,
451 topk_weights=topk_weights,
452 topk_ids=topk_ids,
508 - shared_experts_input=shared_input,
453 + shared_experts_input=shared_experts_input,
454 )
455
511 - if isinstance(result, tuple):
512 - assert shared_output is None
513 - shared_output, hidden_states = result
514 - else:
515 - hidden_states = result
516 -
517 - if not run_shared_experts_before and self.has_separate_shared_experts:
518 - assert shared_output is None
519 - shared_output = self._apply_shared_experts(shared_input, True)
456 + self._maybe_apply_shared_experts(
457 + shared_experts_input,
458 + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
459 + )
460
521 - return shared_output, hidden_states
461 + return (
462 + self.shared_experts.output if self.shared_experts is not None else None,
463 + fused_out,
464 + )
465
466 def _sequence_parallel_context(self):
467 ctx = get_forward_context()
@@ -558,18 +501,16 @@ class DefaultMoERunner(MoERunner):
501
502 return final_shared_hidden_states, final_fused_hidden_states
503
561 - def _maybe_gate(
504 + def _maybe_sync_shared_experts_stream(
505 self,
563 - hidden_states: torch.Tensor,
564 - router_logits: torch.Tensor,
565 - ) -> torch.Tensor:
506 + shared_experts_input: torch.Tensor | None,
507 + ):
508 # If router/gate provided, then apply it here.
509 # (Note: This code runs only when "overlapped mode" is on to allow
510 # parallel execution of shared experts with the FusedMoE via
511 # separate cuda stream)
570 - if self.gate is not None:
571 - router_logits, _ = self.gate(hidden_states)
572 - return router_logits
512 + if self.shared_experts is not None:
513 + self.shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
514
515 @property
516 def do_naive_dispatch_combine(self) -> bool:
@@ -624,7 +565,6 @@ class DefaultMoERunner(MoERunner):
565 hidden_states,
566 dim=0,
567 )
627 - # need RS for shared_output?
568
569 if self.shared_experts is not None:
570 assert shared_output is not None
@@ -637,30 +577,86 @@ class DefaultMoERunner(MoERunner):
577 hidden_states: torch.Tensor,
578 router_logits: torch.Tensor,
579 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
640 - # For latent MoE: save ORIGINAL hidden_states before transform
641 - # (shared_experts need original dimension, routed experts use transformed)
642 - if self.shared_experts is not None:
643 - original_hidden_states = hidden_states
644 - else:
645 - original_hidden_states = None
580 + """Invoke the fused moe layer.
581 +
582 + Input:
583 + - hidden_states
584 + - router_logits
585 +
586 + Output:
587 + - The new hidden_states.
588 + or
589 + - A tuple of (shared experts output, new hidden_states).
590 +
591 + Calling sequence
592 + - forward
593 + - self.forward_entry (_moe_forward or _moe_forward_shared custom op)
594 + - forward_dispatch
595 + - forward_impl (_forward_impl or _forward_impl_chunked)
596 +
597 + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
598 + to the following reasons:
599 + 1. the chunking loop in _forward_impl_chunked cannot be compiled by
600 + torch.compile
601 + 2. pytorch cannot handle union types in custom op signatures so _moe_forward
602 + and _moe_forward_shared must be split.
603 +
604 + If _forward_impl_chunked can be implemented via torch.scan we can potentially
605 + get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence
606 + into the 'forward' method.
607 + """
608
609 # Apply transform for routed experts (e.g., latent projection for latent MoE)
648 - hidden_states = self.apply_routed_input_transform(hidden_states)
610 + hidden_states, shared_experts_input = self.apply_routed_input_transform(
611 + hidden_states
612 + )
613
614 hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(
651 - original_hidden_states,
615 + shared_experts_input,
616 hidden_states,
617 )
618
655 - fused_output = self.moe_forward(
619 + fused_output = self.forward_entry(
620 hidden_states,
621 router_logits,
658 - original_hidden_states,
622 + shared_experts_input,
623 self._encode_layer_name(),
624 )
625
626 return self._maybe_reduce_output(fused_output, og_hidden_dims)
627
628 + def forward_dispatch(
629 + self,
630 + layer: torch.nn.Module,
631 + hidden_states: torch.Tensor,
632 + router_logits: torch.Tensor,
633 + shared_experts_input: torch.Tensor | None,
634 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
635 + # TODO(bnell): this can be removed after MK migration is complete.
636 + layer.ensure_moe_quant_config_init()
637 +
638 + # Sync aux and main stream for shared expert multi-stream overlap.
639 + self._maybe_sync_shared_experts_stream(shared_experts_input)
640 +
641 + # If the Runner holds the gate, apply it after the stream sync,
642 + # so it can run overlapped with the
643 + # NOTE: in future PR, MoE runner will always hold the gate.
644 + if self.gate is not None:
645 + router_logits, _ = self.gate(hidden_states)
646 +
647 + self._maybe_apply_shared_experts(
648 + shared_experts_input,
649 + SharedExpertsOrder.EXTERNAL,
650 + )
651 +
652 + with self._sequence_parallel_context():
653 + return self.forward_impl(
654 + layer,
655 + hidden_states,
656 + router_logits,
657 + shared_experts_input,
658 + )
659 +
660 def _slice_and_copy_input(
661 self,
662 out_slice: torch.Tensor,
@@ -681,17 +677,13 @@ class DefaultMoERunner(MoERunner):
677 out_slice.copy_(orig_slice, non_blocking=True)
678 return out_slice
679
684 - def forward_impl_chunked(
680 + def _forward_impl_chunked(
681 self,
682 layer: torch.nn.Module,
683 hidden_states: torch.Tensor,
684 router_logits: torch.Tensor,
689 - shared_input: torch.Tensor | None,
685 + shared_experts_input: torch.Tensor | None,
686 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
691 - # Gate overlap not supported when chunking is enabled. Run the
692 - # gate first.
693 - router_logits = self._maybe_gate(hidden_states, router_logits)
694 -
687 final_shared_hidden_states, final_fused_hidden_states = (
688 self._allocate_dp_chunking_outputs(hidden_states, router_logits)
689 )
@@ -737,9 +729,9 @@ class DefaultMoERunner(MoERunner):
729 chunk_end,
730 )
731
740 - shared_input_chunk = (
741 - shared_input[chunk_start:chunk_end, :]
742 - if shared_input is not None
732 + shared_experts_input_chunk = (
733 + shared_experts_input[chunk_start:chunk_end, :]
734 + if shared_experts_input is not None
735 else None
736 )
737
@@ -747,7 +739,7 @@ class DefaultMoERunner(MoERunner):
739 layer=layer,
740 hidden_states=hidden_states_chunk,
741 router_logits=router_logits_chunk,
750 - shared_input=shared_input_chunk,
742 + shared_experts_input=shared_experts_input_chunk,
743 )
744
745 # Store outputs
@@ -769,40 +761,13 @@ class DefaultMoERunner(MoERunner):
761 assert final_shared_hidden_states is not None
762 return (final_shared_hidden_states, final_fused_hidden_states)
763
772 - def forward_impl(
764 + def _forward_impl(
765 self,
766 layer: torch.nn.Module,
767 hidden_states: torch.Tensor,
768 router_logits: torch.Tensor,
777 - shared_input: torch.Tensor | None,
769 + shared_experts_input: torch.Tensor | None,
770 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
779 - self.use_shared_experts_stream = (
780 - current_platform.is_cuda()
781 - and self.has_separate_shared_experts
782 - and not self.use_dp_chunking
783 - and self.shared_experts_stream is not None
784 - and (
785 - hidden_states.shape[0]
786 - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
787 - )
788 - )
789 -
790 - # Check if we need to run shared experts before matrix multiply because
791 - # matrix multiply may modify the hidden_states.
792 - run_shared_experts_before = (
793 - self.has_separate_shared_experts and not self.use_shared_experts_stream
794 - )
795 -
796 - # The shared experts stream must be set up before calling the gate so they
797 - # can be overlapped.
798 - if not run_shared_experts_before:
799 - self._maybe_setup_shared_experts_stream(
800 - hidden_states,
801 - shared_input,
802 - )
803 -
804 - router_logits = self._maybe_gate(hidden_states, router_logits)
805 -
771 # TODO(bnell): parts of the dispatch/combine steps will go away once
772 # #32567 lands and the remaining kernels are made MKs. The PCP
773 # code will probably remain
@@ -816,8 +781,7 @@ class DefaultMoERunner(MoERunner):
781 layer=layer,
782 hidden_states=hidden_states,
783 router_logits=router_logits,
819 - shared_input=shared_input,
820 - run_shared_experts_before=run_shared_experts_before,
784 + shared_experts_input=shared_experts_input,
785 )
786
787 return self._maybe_combine(
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+4
@@ -32,3 +32,7 @@ class MoERunner(ABC):
32 final_hidden_states: torch.Tensor,
33 ):
34 raise NotImplementedError
35 +
36 + @abstractmethod
37 + def is_internal_router(self) -> bool:
38 + raise NotImplementedError
vllm/model_executor/layers/fused_moe/runner/shared_experts.py new
+216
@@ -0,0 +1,216 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from enum import IntEnum
4 +
5 +import torch
6 +
7 +import vllm.envs as envs
8 +from vllm.distributed import (
9 + get_tensor_model_parallel_world_size,
10 + tensor_model_parallel_all_reduce,
11 +)
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FusedMoEConfig,
15 +)
16 +from vllm.model_executor.layers.quantization.base_config import (
17 + QuantizeMethodBase,
18 +)
19 +from vllm.platforms import current_platform
20 +from vllm.utils.torch_utils import (
21 + aux_stream,
22 + current_stream,
23 +)
24 +from vllm.v1.worker.ubatching import (
25 + dbo_current_ubatch_id,
26 +)
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class SharedExpertsOrder(IntEnum):
32 + # No shared experts.
33 + NONE = (0,)
34 +
35 + # Get rid of this one? combine with BEFORE?
36 + # Note: this might be important for torch.compile reasons. Can
37 + # get rid of it after _moe_forward is undone.
38 + EXTERNAL = (1,)
39 +
40 + # No overlap - defensively called before MK.
41 + NO_OVERLAP = (2,)
42 +
43 + # Overlapped with dispatch/combine in DP/EP - called by the MK.
44 + MK_INTERNAL_OVERLAPPED = (3,)
45 +
46 + # Overlapped with the gate, router, experts in aux stream.
47 + MULTI_STREAM_OVERLAPPED = (4,)
48 +
49 +
50 +class SharedExperts:
51 + def __init__(
52 + self,
53 + layer: torch.nn.Module,
54 + moe_config: FusedMoEConfig,
55 + quant_method: QuantizeMethodBase,
56 + reduce_results: bool,
57 + enable_dbo: bool,
58 + ):
59 + from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
60 + FusedMoEMethodBase,
61 + )
62 +
63 + # quant_method must be a FusedMoEMethodBase but we can't use the type
64 + # due to circular imports.
65 + assert isinstance(quant_method, FusedMoEMethodBase)
66 +
67 + # The SharedExperts need to handle DBO since they can be called from
68 + # an MK's finalize method. We keep a list of outputs indexed by current
69 + # DBO ubatch id to handle this case. If DBO is not enabled, the
70 + # index is always 0 and the second output list element is ignored.
71 + self.enable_dbo = enable_dbo
72 + self._output: list[torch.Tensor | None] = [None, None]
73 + self._layer = layer
74 + self._moe_config = moe_config
75 + self._quant_method = quant_method
76 + self._reduce_results = reduce_results
77 + self._use_dp_chunking = moe_config.moe_parallel_config.use_dp_chunking
78 +
79 + # Allow disabling of the separate shared experts stream for
80 + # debug purposes.
81 + # TODO: Remove this after more extensive testings with TP/DP
82 + # and other execution modes
83 + if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:
84 + logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local")
85 + self._stream = None
86 + else:
87 + # TODO(rob): enable shared expert overlap with non-cuda-alike.
88 + # aux_stream() returns None on non-cuda-alike platforms.
89 + self._stream = aux_stream()
90 + if self._stream is not None:
91 + logger.debug_once(
92 + "Enabled separate cuda stream for MoE shared_experts", scope="local"
93 + )
94 +
95 + @property
96 + def _has_external_experts(self) -> bool:
97 + # Disable shared expert overlap if:
98 + # - we are using eplb with non-default backend, because of correctness issues
99 + # - we are using flashinfer with DP, since there nothing to gain
100 + backend = self._moe_config.moe_parallel_config.all2all_backend
101 + return not (
102 + (
103 + self._moe_config.moe_parallel_config.enable_eplb
104 + and backend != "allgather_reducescatter"
105 + )
106 + or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels
107 + )
108 +
109 + def _determine_shared_experts_order(
110 + self,
111 + hidden_states: torch.Tensor,
112 + ) -> SharedExpertsOrder:
113 + if self._has_external_experts and not self._use_dp_chunking:
114 + return SharedExpertsOrder.EXTERNAL
115 +
116 + if self._quant_method.mk_owns_shared_expert:
117 + return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
118 +
119 + should_run_shared_in_aux_stream = (
120 + current_platform.is_cuda()
121 + and not self._use_dp_chunking
122 + and self._stream is not None
123 + and hidden_states.shape[0]
124 + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
125 + )
126 +
127 + if should_run_shared_in_aux_stream:
128 + return SharedExpertsOrder.MULTI_STREAM_OVERLAPPED
129 + else:
130 + return SharedExpertsOrder.NO_OVERLAP
131 +
132 + def maybe_sync_shared_experts_stream(
133 + self,
134 + shared_experts_input: torch.Tensor,
135 + ):
136 + experts_order = self._determine_shared_experts_order(shared_experts_input)
137 +
138 + if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
139 + assert self._stream is not None
140 + assert self._moe_config.disable_inplace
141 +
142 + # Record that the clone will be used by shared_experts_stream
143 + # to avoid gc issue from deallocation of hidden_states_clone
144 + # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
145 + # NOTE: We don't need shared_output.record_stream(current_stream())
146 + # because we synch the streams before using shared_output.
147 + shared_experts_input.record_stream(self._stream)
148 +
149 + # Mark sync start point for the aux stream since we will
150 + # run in parallel with router/gate.
151 + self._stream.wait_stream(current_stream())
152 +
153 + def _run_in_aux_stream(
154 + self,
155 + shared_experts_input: torch.Tensor,
156 + ) -> torch.Tensor:
157 + # TODO: assert that maybe_sync_shared_experts_stream has been called.
158 +
159 + # Run shared experts in parallel on a separate stream.
160 + with torch.cuda.stream(self._stream):
161 + output = self._layer(shared_experts_input)
162 + current_stream().wait_stream(self._stream)
163 +
164 + return output
165 +
166 + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor:
167 + # Reduce shared expert outputs if necessary, since the MLP
168 + # should have been created with reduce_results=False.
169 + if (
170 + self._reduce_results
171 + and self._quant_method.moe_kernel is not None
172 + and self._quant_method.moe_kernel.output_is_reduced()
173 + and get_tensor_model_parallel_world_size() > 1
174 + ):
175 + shared_out = tensor_model_parallel_all_reduce(shared_out)
176 + return shared_out
177 +
178 + @property
179 + def _output_idx(self) -> int:
180 + return dbo_current_ubatch_id() if self.enable_dbo else 0
181 +
182 + @property
183 + def output(self) -> torch.Tensor:
184 + assert self._output[self._output_idx] is not None
185 + output = self._output[self._output_idx]
186 + self._output[self._output_idx] = None
187 + return output
188 +
189 + def apply(
190 + self,
191 + shared_experts_input: torch.Tensor,
192 + order: SharedExpertsOrder,
193 + ):
194 + experts_order = self._determine_shared_experts_order(shared_experts_input)
195 +
196 + if order != experts_order:
197 + return None
198 +
199 + assert self._output[self._output_idx] is None
200 +
201 + if order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
202 + self._output[self._output_idx] = self._run_in_aux_stream(
203 + shared_experts_input
204 + )
205 + else:
206 + self._output[self._output_idx] = self._layer(shared_experts_input)
207 +
208 + if order == SharedExpertsOrder.EXTERNAL:
209 + # TODO: figure out how to combine this with maybe_reduce_output?
210 + # or get rid of it completely.
211 + assert self._output[self._output_idx] is not None
212 + self._output[self._output_idx] = self._maybe_reduce_shared_out(
213 + self._output[self._output_idx]
214 + )
215 +
216 + assert self._output[self._output_idx] is not None
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
+8 -37
@@ -3,14 +3,10 @@
3
4 import torch
5
6 -from vllm.distributed import (
7 - get_tensor_model_parallel_world_size,
8 - tensor_model_parallel_all_reduce,
9 -)
6 from vllm.model_executor.layers.fused_moe.layer import FusedMoE
7
8
13 -# TODO(bnell): Add shared + fused combo function? e.g. +
9 +# TODO(bnell): Remove this entirely
10 class SharedFusedMoE(FusedMoE):
11 """
12 A FusedMoE operation that also computes the results of shared experts.
@@ -23,36 +19,11 @@ class SharedFusedMoE(FusedMoE):
19 hidden_states: torch.Tensor,
20 router_logits: torch.Tensor,
21 ) -> tuple[torch.Tensor, torch.Tensor]:
26 - if not self.use_overlapped:
27 - if self._shared_experts is not None:
28 - shared_out = self._shared_experts(hidden_states)
29 -
30 - # Reduce shared expert outputs if necessary, since the MLP
31 - # should have been created with reduce_results=False.
32 - if (
33 - self.reduce_results
34 - and get_tensor_model_parallel_world_size() > 1
35 - and self.must_reduce_shared_expert_outputs()
36 - ):
37 - shared_out = tensor_model_parallel_all_reduce(shared_out)
38 - else:
39 - shared_out = None
40 -
41 - fused_out = super().forward(
42 - hidden_states=hidden_states,
43 - router_logits=router_logits,
44 - )
22 + result = super().forward(
23 + hidden_states=hidden_states,
24 + router_logits=router_logits,
25 + )
26 + if self.shared_experts is None:
27 + return None, result
28 else:
46 - shared_out, fused_out = super().forward(
47 - hidden_states=hidden_states,
48 - router_logits=router_logits,
49 - )
50 - # ensure early TP reduction of shared expert outputs when required
51 - if (
52 - shared_out is not None
53 - and self.reduce_results
54 - and get_tensor_model_parallel_world_size() > 1
55 - and self.must_reduce_shared_expert_outputs()
56 - ):
57 - shared_out = tensor_model_parallel_all_reduce(shared_out)
58 - return shared_out, fused_out
29 + return result
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+4 -4
@@ -245,7 +245,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
245 topk_weights: torch.Tensor,
246 topk_ids: torch.Tensor,
247 shared_experts_input: torch.Tensor | None,
248 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
248 + ) -> torch.Tensor:
249 return self.forward(
250 layer=layer,
251 x=x,
@@ -261,7 +261,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
261 topk_weights: torch.Tensor,
262 topk_ids: torch.Tensor,
263 shared_experts_input: torch.Tensor | None,
264 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
264 + ) -> torch.Tensor:
265 assert self.moe_kernel is not None
266 return self.moe_kernel.apply(
267 hidden_states=x,
@@ -283,7 +283,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
283 topk_weights: torch.Tensor,
284 topk_ids: torch.Tensor,
285 shared_experts_input: torch.Tensor | None,
286 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
286 + ) -> torch.Tensor:
287 return self.forward_native(
288 layer, x, topk_weights, topk_ids, shared_experts_input
289 )
@@ -293,7 +293,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
293 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
294 x: torch.Tensor,
295 router_logits: torch.Tensor,
296 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
296 + ) -> torch.Tensor:
297 assert self.is_monolithic
298 if self.unquantized_backend == UnquantizedMoeBackend.CPU:
299 assert self.moe_kernel is None
vllm/model_executor/layers/quantization/awq_marlin.py
+1 -1
@@ -811,7 +811,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
811 topk_weights: torch.Tensor,
812 topk_ids: torch.Tensor,
813 shared_experts_input: torch.Tensor | None,
814 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
814 + ) -> torch.Tensor:
815 return fused_marlin_moe(
816 x,
817 layer.w13_qweight,
vllm/model_executor/layers/quantization/bitsandbytes.py
+1 -1
@@ -483,7 +483,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
483 topk_weights: torch.Tensor,
484 topk_ids: torch.Tensor,
485 shared_experts_input: torch.Tensor | None,
486 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
486 + ) -> torch.Tensor:
487 from vllm.model_executor.layers.fused_moe import fused_experts
488
489 # TODO(bnell): Do these need to be called on the hot path?
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+10 -10
@@ -355,7 +355,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
355 topk_weights: torch.Tensor,
356 topk_ids: torch.Tensor,
357 shared_experts_input: torch.Tensor | None,
358 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
358 + ) -> torch.Tensor:
359 assert self.moe_kernel is not None
360 return self.moe_kernel.apply(
361 x,
@@ -603,7 +603,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
603 layer: FusedMoE,
604 x: torch.Tensor,
605 router_logits: torch.Tensor,
606 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
606 + ) -> torch.Tensor:
607 assert self.is_monolithic
608 assert self.moe_kernel is not None
609 return self.moe_kernel.apply_monolithic(
@@ -628,7 +628,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
628 topk_weights: torch.Tensor,
629 topk_ids: torch.Tensor,
630 shared_experts_input: torch.Tensor | None,
631 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
631 + ) -> torch.Tensor:
632 assert self.moe_kernel is not None
633 return self.moe_kernel.apply(
634 x,
@@ -963,7 +963,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
963 layer: FusedMoE,
964 x: torch.Tensor,
965 router_logits: torch.Tensor,
966 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
966 + ) -> torch.Tensor:
967 assert self.moe_kernel is not None
968 return self.moe_kernel.apply_monolithic(
969 x,
@@ -987,7 +987,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
987 topk_weights: torch.Tensor,
988 topk_ids: torch.Tensor,
989 shared_experts_input: torch.Tensor | None,
990 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
990 + ) -> torch.Tensor:
991 assert not self.is_monolithic
992 assert self.moe_kernel is not None
993 return self.moe_kernel.apply(
@@ -1127,7 +1127,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1127 topk_weights: torch.Tensor,
1128 topk_ids: torch.Tensor,
1129 shared_experts_input: torch.Tensor | None,
1130 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1130 + ) -> torch.Tensor:
1131 from vllm.model_executor.layers.fused_moe import fused_experts
1132
1133 return fused_experts(
@@ -1611,7 +1611,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1611 layer: FusedMoE,
1612 x: torch.Tensor,
1613 router_logits: torch.Tensor,
1614 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1614 + ) -> torch.Tensor:
1615 assert self.kernel_backend == "Flashinfer"
1616 return flashinfer_trtllm_mxint4_moe(
1617 x=x,
@@ -1638,7 +1638,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1638 topk_weights: torch.Tensor,
1639 topk_ids: torch.Tensor,
1640 shared_experts_input: torch.Tensor | None,
1641 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1641 + ) -> torch.Tensor:
1642 assert self.kernel_backend == "Marlin"
1643 return fused_marlin_moe(
1644 x,
@@ -1887,7 +1887,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1887 topk_weights: torch.Tensor,
1888 topk_ids: torch.Tensor,
1889 shared_experts_input: torch.Tensor | None,
1890 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1890 + ) -> torch.Tensor:
1891 from vllm.model_executor.layers.fused_moe import fused_experts
1892
1893 return fused_experts(
@@ -2502,7 +2502,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2502 topk_weights: torch.Tensor,
2503 topk_ids: torch.Tensor,
2504 shared_experts_input: torch.Tensor | None,
2505 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2505 + ) -> torch.Tensor:
2506 if layer.enable_eplb:
2507 raise NotImplementedError(
2508 "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
vllm/model_executor/layers/quantization/experts_int8.py
+1 -1
@@ -141,7 +141,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
141 topk_weights: torch.Tensor,
142 topk_ids: torch.Tensor,
143 shared_experts_input: torch.Tensor | None,
144 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
144 + ) -> torch.Tensor:
145 from vllm.model_executor.layers.fused_moe import fused_experts
146
147 return fused_experts(
vllm/model_executor/layers/quantization/fp8.py
+2 -2
@@ -877,7 +877,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
877 layer: FusedMoE,
878 x: torch.Tensor,
879 router_logits: torch.Tensor,
880 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
880 + ) -> torch.Tensor:
881 assert self.is_monolithic
882 assert self.moe_kernel is not None
883 return self.moe_kernel.apply_monolithic(
@@ -902,7 +902,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
902 topk_weights: torch.Tensor,
903 topk_ids: torch.Tensor,
904 shared_experts_input: torch.Tensor | None,
905 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
905 + ) -> torch.Tensor:
906 assert not self.is_monolithic
907 assert self.moe_kernel is not None
908 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/gguf.py
+1 -1
@@ -650,7 +650,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):
650 topk_weights: torch.Tensor,
651 topk_ids: torch.Tensor,
652 shared_experts_input: torch.Tensor | None,
653 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
653 + ) -> torch.Tensor:
654 if layer.apply_router_weight_on_input:
655 raise NotImplementedError(
656 "Apply router weight on input is not supported for"
vllm/model_executor/layers/quantization/gptq_marlin.py
+1 -1
@@ -907,7 +907,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
907 topk_weights: torch.Tensor,
908 topk_ids: torch.Tensor,
909 shared_experts_input: torch.Tensor | None,
910 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
910 + ) -> torch.Tensor:
911 return fused_marlin_moe(
912 x,
913 layer.w13_qweight,
vllm/model_executor/layers/quantization/modelopt.py
+4 -4
@@ -935,7 +935,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
935 layer: FusedMoE,
936 x: torch.Tensor,
937 router_logits: torch.Tensor,
938 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
938 + ) -> torch.Tensor:
939 assert self.is_monolithic
940 assert self.moe_kernel is not None
941 return self.moe_kernel.apply_monolithic(
@@ -960,7 +960,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
960 topk_weights: torch.Tensor,
961 topk_ids: torch.Tensor,
962 shared_experts_input: torch.Tensor | None,
963 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
963 + ) -> torch.Tensor:
964 assert not self.is_monolithic
965 assert self.moe_kernel is not None
966 return self.moe_kernel.apply(
@@ -1419,7 +1419,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1419 layer: FusedMoE,
1420 x: torch.Tensor,
1421 router_logits: torch.Tensor,
1422 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1422 + ) -> torch.Tensor:
1423 assert self.is_monolithic
1424 assert self.moe_kernel is not None
1425 return self.moe_kernel.apply_monolithic(
@@ -1444,7 +1444,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1444 topk_weights: torch.Tensor,
1445 topk_ids: torch.Tensor,
1446 shared_experts_input: torch.Tensor | None,
1447 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1447 + ) -> torch.Tensor:
1448 assert not self.is_monolithic
1449 assert self.moe_kernel is not None
1450 return self.moe_kernel.apply(
vllm/model_executor/layers/quantization/moe_wna16.py
+1 -1
@@ -369,7 +369,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
369 topk_weights: torch.Tensor,
370 topk_ids: torch.Tensor,
371 shared_experts_input: torch.Tensor | None,
372 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
372 + ) -> torch.Tensor:
373 from vllm.model_executor.layers.fused_moe import fused_experts
374
375 assert layer.activation == MoEActivation.SILU, (
vllm/model_executor/layers/quantization/mxfp4.py
+2 -2
@@ -377,7 +377,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
377 topk_weights: torch.Tensor,
378 topk_ids: torch.Tensor,
379 shared_experts_input: torch.Tensor | None,
380 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
380 + ) -> torch.Tensor:
381 assert not self.is_monolithic
382 assert self.moe_kernel is not None
383 return self.moe_kernel.apply(
@@ -398,7 +398,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
398 layer: FusedMoE,
399 x: torch.Tensor,
400 router_logits: torch.Tensor,
401 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
401 + ) -> torch.Tensor:
402 assert self.is_monolithic
403 assert self.moe_kernel is not None
404 return self.moe_kernel.apply_monolithic(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3 -3
@@ -444,7 +444,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
444 topk_weights: torch.Tensor,
445 topk_ids: torch.Tensor,
446 shared_experts_input: torch.Tensor | None,
447 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
447 + ) -> torch.Tensor:
448 if self.rocm_aiter_moe_enabled:
449 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
450 rocm_aiter_fused_experts,
@@ -634,7 +634,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
634 topk_weights: torch.Tensor,
635 topk_ids: torch.Tensor,
636 shared_experts_input: torch.Tensor | None,
637 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
637 + ) -> torch.Tensor:
638 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
639 rocm_aiter_fused_experts,
640 )
@@ -1027,7 +1027,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1027 topk_weights: torch.Tensor,
1028 topk_ids: torch.Tensor,
1029 shared_experts_input: torch.Tensor | None,
1030 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1030 + ) -> torch.Tensor:
1031 if not self.emulate:
1032 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1033 rocm_aiter_fused_experts,
vllm/model_executor/models/transformers/moe.py
+2
@@ -94,6 +94,8 @@ def transformers_moe_forward(
94 self = forward_context.no_compile_layers[layer_name]
95 self._topk_ids = topk_ids
96 # Clone hidden_states because it will be mutated in-place in FusedMoE
97 + # TODO(bnell): figure out a way to avoid calling runner directly.
98 + # it is a hack that the weight are being passed via logits.
99 return self.runner.forward(hidden_states.clone(), topk_weights)
100
101