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, 11:15 UTC
[MoE] Make MoERunnerInterface a PluggableLayer for OOT support (#35178)
Signed-off-by: wxsIcey <1790571317@qq.com> Signed-off-by: Icey <1790571317@qq.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Xiaoshuang Wang committed
Apr 30, 2026 at 18:31 UTC
b55b26520c088eeed791cbe4648c73ea015f3613
4 files changed
+35
-23
docs/mkdocs/hooks/generate_argparse.py
+13
-1
@@ -38,8 +38,20 @@ class MockCustomOp:
38
return decorator
39
40
41
+class MockPluggableLayer:
42
+ @staticmethod
43
+ def register(name):
44
+ def decorator(cls):
45
+ return cls
46
+
47
+ return decorator
48
+
49
+
50
mock_if_no_torch("vllm._C", MagicMock())
42
-mock_if_no_torch("vllm.model_executor.custom_op", MagicMock(CustomOp=MockCustomOp))
51
+mock_if_no_torch(
52
+ "vllm.model_executor.custom_op",
53
+ MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
54
+)
55
mock_if_no_torch(
56
"vllm.utils.torch_utils", MagicMock(direct_register_custom_op=lambda *a, **k: None)
57
)
vllm/model_executor/layers/fused_moe/layer.py
+11
-12
@@ -1487,15 +1487,19 @@ class FusedMoE(PluggableLayer):
1487
"w2_input_scale",
1488
}
1489
1490
+ # Parameters of non-expert submodules that live inside runner (MoERunner).
1491
+ # These must be excluded from EPLB weight rearrangement.
1492
+ NON_EXPERT_PREFIXES = (
1493
+ "runner._shared_experts.",
1494
+ "runner.gate.",
1495
+ "runner.routed_input_transform.",
1496
+ "runner.routed_output_transform.",
1497
+ )
1498
+
1499
assert all(
1500
weight.is_contiguous()
1501
for name, weight in weights
1493
- if not (
1494
- name.startswith("_shared_experts.")
1495
- or name.startswith("_gate.")
1496
- or name.startswith("_routed_input_transform.")
1497
- or name.startswith("_routed_output_transform.")
1498
- )
1502
+ if not name.startswith(NON_EXPERT_PREFIXES)
1503
and name not in NON_EXPERT_WEIGHTS
1504
)
1505
@@ -1504,12 +1508,7 @@ class FusedMoE(PluggableLayer):
1508
for name, weight in weights
1509
if name not in NON_EXPERT_WEIGHTS
1510
and weight.shape != torch.Size([])
1507
- and not name.startswith("_shared_experts.")
1508
- # exclude parameters from non-expert submodules,
1509
- # e.g. gate/shared/transforms.
1510
- and not name.startswith("_gate.")
1511
- and not name.startswith("_routed_input_transform.")
1512
- and not name.startswith("_routed_output_transform.")
1511
+ and not name.startswith(NON_EXPERT_PREFIXES)
1512
]
1513
1514
def set_eplb_state(
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+9
-9
@@ -220,7 +220,7 @@ class MoERunner(MoERunnerInterface):
220
self.routed_output_transform = routed_output_transform
221
self.routed_scaling_factor = routed_scaling_factor
222
self.gate = gate
223
- self.quant_method = quant_method
223
+ self._quant_method = quant_method
224
self.enable_dbo = enable_dbo
225
226
self._shared_experts: SharedExperts | None = None
@@ -263,7 +263,7 @@ class MoERunner(MoERunnerInterface):
263
def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
264
if self._shared_experts is not None:
265
self._shared_experts._quant_method = quant_method
266
- self.quant_method = quant_method
266
+ self._quant_method = quant_method
267
268
def is_internal_router(self) -> bool:
269
return self.gate is not None
@@ -330,8 +330,8 @@ class MoERunner(MoERunnerInterface):
330
@property
331
def _fused_output_is_reduced(self) -> bool:
332
return (
333
- self.quant_method.moe_kernel is not None
334
- and self.quant_method.moe_kernel.output_is_reduced()
333
+ self._quant_method.moe_kernel is not None
334
+ and self._quant_method.moe_kernel.output_is_reduced()
335
)
336
337
def _maybe_reduce_shared_expert_output(
@@ -407,7 +407,7 @@ class MoERunner(MoERunnerInterface):
407
)
408
transformed_hidden_dim = hidden_states.shape[-1]
409
if (
410
- not self.quant_method.skip_forward_padding
410
+ not self._quant_method.skip_forward_padding
411
and self.moe_config.hidden_dim != transformed_hidden_dim
412
):
413
hidden_states = F.pad(
@@ -451,8 +451,8 @@ class MoERunner(MoERunnerInterface):
451
shared_experts_input, SharedExpertsOrder.NO_OVERLAP
452
)
453
454
- if self.quant_method.is_monolithic:
455
- fused_out = self.quant_method.apply_monolithic(
454
+ if self._quant_method.is_monolithic:
455
+ fused_out = self._quant_method.apply_monolithic(
456
layer=layer,
457
x=hidden_states,
458
router_logits=router_logits,
@@ -467,7 +467,7 @@ class MoERunner(MoERunnerInterface):
467
468
# Passing shared_experts_input in case SharedExpertsOrder is
469
# MK_INTERNAL_OVERLAPPED.
470
- fused_out = self.quant_method.apply(
470
+ fused_out = self._quant_method.apply(
471
layer=layer,
472
x=hidden_states,
473
topk_weights=topk_weights,
@@ -618,7 +618,7 @@ class MoERunner(MoERunnerInterface):
618
@property
619
def do_naive_dispatch_combine(self) -> bool:
620
return (
621
- self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk
621
+ self.moe_config.dp_size > 1 and not self._quant_method.supports_internal_mk
622
)
623
624
def _maybe_dispatch(
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
+2
-1
@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
4
5
import torch
6
7
+from vllm.model_executor.custom_op import PluggableLayer
8
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
9
FusedMoEMethodBase,
10
)
@@ -12,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
13
)
14
15
15
-class MoERunnerInterface(ABC):
16
+class MoERunnerInterface(PluggableLayer, ABC):
17
"""
18
Abstract base class for Mixture of Experts (MoE) runners.
19