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, 04:45 UTC
[PluggableLayer][3/N] Apply PluggableLayer to moe-related layers. (#33556)
Signed-off-by: whx-sjtu <2952154980@qq.com>
Hexiang Wang committed
Apr 14, 2026 at 21:55 UTC
f02b3269e7e0ebfdc62f39db2169d88074512d6d
2 files changed
+6
-13
vllm/model_executor/layers/fused_moe/layer.py
+4
-11
@@ -18,7 +18,7 @@ from vllm.distributed import (
18
)
19
from vllm.distributed.eplb.eplb_state import EplbLayerState, EplbState
20
from vllm.logger import init_logger
21
-from vllm.model_executor.custom_op import CustomOp
21
+from vllm.model_executor.custom_op import PluggableLayer
22
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
23
from vllm.model_executor.layers.fused_moe.config import (
24
FusedMoEConfig,
@@ -213,8 +213,8 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
213
214
215
# --8<-- [start:fused_moe]
216
-@CustomOp.register("fused_moe")
217
-class FusedMoE(CustomOp):
216
+@PluggableLayer.register("fused_moe")
217
+class FusedMoE(PluggableLayer):
218
"""FusedMoE layer for MoE models.
219
220
This layer contains both MergedColumnParallel weights (gate_up_proj /
@@ -1532,7 +1532,7 @@ class FusedMoE(CustomOp):
1532
"""
1533
return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states)
1534
1535
- def forward_native(
1535
+ def forward(
1536
self,
1537
hidden_states: torch.Tensor,
1538
router_logits: torch.Tensor,
@@ -1548,13 +1548,6 @@ class FusedMoE(CustomOp):
1548
self._expert_map if not self.rocm_aiter_fmoe_enabled else self.expert_mask
1549
)
1550
1551
- def forward_cuda(
1552
- self,
1553
- hidden_states: torch.Tensor,
1554
- router_logits: torch.Tensor,
1555
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1556
- return self.forward_native(hidden_states, router_logits)
1557
-
1551
@classmethod
1552
def make_expert_params_mapping(
1553
cls,
vllm/model_executor/models/transformers/moe.py
+2
-2
@@ -24,7 +24,7 @@ import torch.nn as nn
24
from vllm.config.utils import getattr_iter
25
from vllm.distributed import get_dp_group, get_ep_group
26
from vllm.forward_context import ForwardContext, get_forward_context
27
-from vllm.model_executor.custom_op import CustomOp
27
+from vllm.model_executor.custom_op import PluggableLayer
28
from vllm.model_executor.layers.fused_moe import FusedMoE
29
from vllm.model_executor.models.interfaces import MixtureOfExperts
30
from vllm.model_executor.models.utils import maybe_prefix
@@ -38,7 +38,7 @@ if TYPE_CHECKING:
38
39
40
# --8<-- [start:transformers_fused_moe]
41
-@CustomOp.register("transformers_fused_moe")
41
+@PluggableLayer.register("transformers_fused_moe")
42
class TransformersFusedMoE(FusedMoE):
43
"""Custom FusedMoE for the Transformers modeling backend."""
44