Fix Transformers backend FP8 MoE and remove some boilerplate (#46820)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

vllm/model_executor/models/AXK1.py

@@ -1088,8 +1088,6 @@ class AXK1ForCausalLM(
         self.set_moe_parameters()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.num_expert_groups = getattr(self.config, "n_group", 1)
 
         self.moe_layers = []

vllm/model_executor/models/afmoe.py

@@ -42,6 +42,7 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 from vllm.model_executor.models.interfaces import (
     EagleModelMixin,
+    MixtureOfExperts,
     SupportsEagle3,
     SupportsLoRA,
     SupportsPP,
@@ -595,7 +596,9 @@ class AfmoeModel(nn.Module, EagleModelMixin):
         return loaded_params
 
 
-class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
+class AfmoeForCausalLM(
+    nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA, MixtureOfExperts
+):
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -635,8 +638,6 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
         self.make_empty_intermediate_tensors = (
             self.model.make_empty_intermediate_tensors
         )
-        self.expert_weights = []
-
         # Set MoE hyperparameters
         self.num_moe_layers = config.num_hidden_layers - config.num_dense_layers
         self.num_expert_groups = config.n_group
@@ -663,21 +664,24 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
             self.num_shared_experts = example_moe.n_shared_experts
             self.num_redundant_experts = example_moe.n_redundant_experts
 
-    def set_eplb_state(
+    def update_physical_experts_metadata(
         self,
-        expert_load_view: torch.Tensor,
-        logical_to_physical_map: torch.Tensor,
-        logical_replica_count: torch.Tensor,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
     ) -> None:
-        for layer_idx, layer in enumerate(self.moe_layers):
-            # Register the expert weights.
-            self.expert_weights.append(layer.get_expert_weights())
-            layer.set_eplb_state(
-                moe_layer_idx=layer_idx,
-                expert_load_view=expert_load_view,
-                logical_to_physical_map=logical_to_physical_map,
-                logical_replica_count=logical_replica_count,
-            )
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+            if layer.moe_enabled:
+                moe = layer.mlp
+                moe.n_local_physical_experts = num_local_physical_experts
+                moe.n_physical_experts = num_physical_experts
+                moe.n_redundant_experts = self.num_redundant_experts
+                moe.experts.update_expert_map()
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)

vllm/model_executor/models/deepseek_mtp.py

@@ -217,7 +217,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
         self.set_moe_parameters()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
         self.num_moe_layers = self.config.num_nextn_predict_layers
         self.num_expert_groups = self.config.n_group
 

vllm/model_executor/models/deepseek_v2.py

@@ -1684,8 +1684,6 @@ class DeepseekV2ForCausalLM(
         self.set_moe_parameters()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.num_expert_groups = getattr(self.config, "n_group", 1)
 
         self.moe_layers = []

vllm/model_executor/models/ernie45_moe.py

@@ -656,8 +656,6 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe
             self.model.make_empty_intermediate_tensors
         )
 
-        self.expert_weights = []
-
         # Set MoE hyperparameters
         moe_layers_indices = [
             i

vllm/model_executor/models/gemma4.py

@@ -1562,7 +1562,6 @@ class Gemma4ForCausalLM(
         )
 
         # --- MixtureOfExperts protocol ---
-        self.expert_weights: list[list[torch.Tensor]] = []
         self.moe_layers: list[nn.Module] = []
         example_moe: Gemma4MoE | None = None
 

vllm/model_executor/models/gemma4_mm.py

@@ -1114,7 +1114,6 @@ class Gemma4ForConditionalGeneration(
                 )
 
         # --- MixtureOfExperts delegation to language_model ---
-        self.expert_weights = self.language_model.expert_weights
         self.moe_layers = self.language_model.moe_layers
         self.num_moe_layers = self.language_model.num_moe_layers
         self.num_logical_experts = self.language_model.num_logical_experts

vllm/model_executor/models/glm4_moe.py

@@ -655,8 +655,6 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, Glm4MixtureOfExper
         self.make_empty_intermediate_tensors = (
             self.model.make_empty_intermediate_tensors
         )
-        self.expert_weights = []
-
         # Set MoE hyperparameters
         self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
         self.num_expert_groups = config.n_group

vllm/model_executor/models/glm4_moe_lite.py

@@ -573,8 +573,6 @@ class Glm4MoeLiteForCausalLM(
         self.set_moe_parameters()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.num_expert_groups = getattr(self.config, "n_group", 1)
 
         self.moe_layers = []

vllm/model_executor/models/glm4_moe_lite_mtp.py

@@ -209,8 +209,6 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
 
-        self.expert_weights = []
-
         # Set MoE hyperparameters
         self.num_moe_layers = self.config.num_nextn_predict_layers
         self.num_expert_groups = self.config.n_group

vllm/model_executor/models/glm4_moe_mtp.py

@@ -195,8 +195,6 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
 
-        self.expert_weights = []
-
         # Set MoE hyperparameters
         self.num_moe_layers = self.config.num_nextn_predict_layers
         self.num_expert_groups = self.config.n_group

vllm/model_executor/models/glm_ocr_mtp.py

@@ -134,7 +134,6 @@ class GlmOcrMTP(nn.Module, SupportsPP):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
 
-        self.expert_weights = []
         self.num_layers = self.config.num_nextn_predict_layers
         for layer in self.model.layers.values():
             assert isinstance(layer, GlmOcrMultiTokenPredictorLayer)

vllm/model_executor/models/hunyuan_v1.py

@@ -991,7 +991,6 @@ class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
 
         # Set MoE hyperparameters
-        self.expert_weights = []
         self.num_expert_groups = 1
         self.moe_layers = []
         example_layer = None

vllm/model_executor/models/hy_v3.py

@@ -68,7 +68,7 @@ from vllm.model_executor.model_loader.weight_utils import (
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.hy_v3 import HYV3Config
 
-from .interfaces import SupportsLoRA, SupportsPP
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -392,7 +392,7 @@ class HYV3DecoderLayer(nn.Module):
 
 
 @support_torch_compile
-class HYV3Model(nn.Module):
+class HYV3Model(nn.Module, MixtureOfExperts):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -429,7 +429,6 @@ class HYV3Model(nn.Module):
         )
 
         # Set MoE hyperparameters
-        self.expert_weights = []
         self.num_expert_groups = 1
         self.moe_layers = []
         example_layer = None

vllm/model_executor/models/interfaces.py

@@ -29,36 +29,34 @@ from torch import Tensor
 from transformers.models.whisper.tokenization_whisper import LANGUAGES
 from typing_extensions import Self, TypeIs
 
-from vllm.config import ModelConfig, SpeechToTextConfig, SpeechToTextParams
-from vllm.inputs import PromptType, TokensPrompt
 from vllm.logger import init_logger
-from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.tasks import ScoreType
 from vllm.utils.collection_utils import common_prefix
 from vllm.utils.func_utils import supports_kw
 
-from .interfaces_base import VllmModel
-
 if TYPE_CHECKING:
-    from vllm.config import VllmConfig
+    from vllm.config import (
+        ModelConfig,
+        SpeechToTextConfig,
+        SpeechToTextParams,
+        VllmConfig,
+    )
+    from vllm.inputs import PromptType, TokensPrompt
     from vllm.lora.model_manager import LoRAModelManager
+    from vllm.model_executor.layers.fused_moe import MoERunner
+    from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
+    from vllm.model_executor.models.interfaces_base import VllmModel
     from vllm.model_executor.models.utils import WeightsMapper
     from vllm.multimodal.inputs import MultiModalFeatureSpec
     from vllm.multimodal.registry import _ProcessorFactories
     from vllm.sequence import IntermediateTensors
+    from vllm.tasks import ScoreType
     from vllm.v1.worker.encoder_cudagraph_defs import (
         EncoderCudaGraphCaptureInputs,
         EncoderCudaGraphConfig,
         EncoderCudaGraphReplayBuffers,
         EncoderItemSpec,
     )
-else:
-    VllmConfig = object
-    WeightsMapper = object
-    MultiModalFeatureSpec = object
-    _ProcessorFactories = object
-    IntermediateTensors = object
 
 logger = init_logger(__name__)
 
@@ -89,7 +87,7 @@ def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:
 
 
 # Cache results of `SupportsMultiModal.get_language_model`
-_language_model_by_module = dict[nn.Module, VllmModel]()
+_language_model_by_module = dict[nn.Module, "VllmModel"]()
 
 
 @runtime_checkable
@@ -123,7 +121,7 @@ class SupportsMultiModal(Protocol):
     in their raw form and not input embeddings.
     """
 
-    _processor_factory: ClassVar[_ProcessorFactories]
+    _processor_factory: ClassVar["_ProcessorFactories"]
     """
     Set internally by `MultiModalRegistry.register_processor`.
     """
@@ -175,7 +173,7 @@ class SupportsMultiModal(Protocol):
             self._has_oov_mm_tokens,
         )
 
-    def get_language_model(self) -> VllmModel:
+    def get_language_model(self) -> "VllmModel":
         """
         Returns the underlying language model used for text generation.
 
@@ -216,7 +214,7 @@ class SupportsMultiModal(Protocol):
     @contextmanager
     def _mark_language_model(
         self,
-        vllm_config: VllmConfig,
+        vllm_config: "VllmConfig",
         *,
         targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
     ):
@@ -251,7 +249,7 @@ class SupportsMultiModal(Protocol):
     @contextmanager
     def _mark_tower_model(
         self,
-        vllm_config: VllmConfig,
+        vllm_config: "VllmConfig",
         modalities: set[str] | str,
         *,
         targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
@@ -295,7 +293,7 @@ class SupportsMultiModal(Protocol):
     @contextmanager
     def _mark_composite_model(
         self,
-        vllm_config: VllmConfig,
+        vllm_config: "VllmConfig",
         *,
         language_targets: type[nn.Module] | tuple[type[nn.Module], ...],
         tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]],
@@ -513,7 +511,7 @@ class SupportsScoreTemplate(Protocol):
         ...
 
     @classmethod
-    def post_process_tokens(cls, prompt: TokensPrompt) -> None:
+    def post_process_tokens(cls, prompt: "TokensPrompt") -> None:
         """
         Perform architecture-specific manipulations on the input tokens.
         """
@@ -633,7 +631,7 @@ class SupportsPP(Protocol):
         batch_size: int,
         dtype: torch.dtype,
         device: torch.device,
-    ) -> IntermediateTensors:
+    ) -> "IntermediateTensors":
         """Called when PP rank > 0 for profiling purposes."""
         ...
 
@@ -642,8 +640,8 @@ class SupportsPP(Protocol):
         input_ids: Tensor | None,
         positions: Tensor,
         *,
-        intermediate_tensors: IntermediateTensors | None,
-    ) -> IntermediateTensors | None:
+        intermediate_tensors: "IntermediateTensors | None",
+    ) -> "IntermediateTensors | None":
         """
         Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
         PP rank > 0.
@@ -665,15 +663,15 @@ class _SupportsPPType(Protocol):
         batch_size: int,
         dtype: torch.dtype,
         device: torch.device,
-    ) -> IntermediateTensors: ...
+    ) -> "IntermediateTensors": ...
 
     def forward(
         self,
         input_ids: Tensor | None,
         positions: Tensor,
         *,
-        intermediate_tensors: IntermediateTensors | None,
-    ) -> Tensor | IntermediateTensors: ...
+        intermediate_tensors: "IntermediateTensors | None",
+    ) -> "Tensor | IntermediateTensors": ...
 
 
 @overload
@@ -803,7 +801,7 @@ class IsHybrid(Protocol):
     @classmethod
     def get_mamba_state_shape_from_config(
         cls,
-        vllm_config: VllmConfig,
+        vllm_config: "VllmConfig",
     ) -> tuple[tuple[int, int], tuple[int, int, int]]:
         """Calculate shapes for Mamba's convolutional and state caches.
 
@@ -818,7 +816,7 @@ class IsHybrid(Protocol):
         ...
 
     @classmethod
-    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:
+    def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]:
         """Calculate copy-function callables for each Mamba state.
 
         Returns:
@@ -883,7 +881,7 @@ class MixtureOfExperts(Protocol):
     num_redundant_experts: int
     """Number of redundant experts in this model."""
 
-    moe_layers: Iterable[nn.Module]
+    moe_layers: Iterable["MoERunner"]
     """List of MoE layers in this model."""
 
     def set_eplb_state(
@@ -908,6 +906,7 @@ class MixtureOfExperts(Protocol):
             logical_to_physical_map: Mapping from logical to physical experts.
             logical_replica_count: Count of replicas for each logical expert.
         """
+        self.expert_weights = []
         for layer_idx, layer in enumerate(self.moe_layers):
             # Register the expert weights.
             self.expert_weights.append(layer.get_expert_weights())
@@ -982,7 +981,7 @@ def supports_mamba_prefix_caching(
 class SupportsCrossEncoding(Protocol):
     """The interface required for all models that support cross encoding."""
 
-    score_type: ClassVar[ScoreType] = "cross-encoder"
+    score_type: ClassVar["ScoreType"] = "cross-encoder"
 
 
 @runtime_checkable
@@ -994,13 +993,13 @@ class SupportsLateInteraction(Protocol):
     MaxSim (max over document tokens, sum over query tokens).
     """
 
-    score_type: ClassVar[ScoreType] = "late-interaction"
+    score_type: ClassVar["ScoreType"] = "late-interaction"
 
 
 class SupportsQuant:
     """The interface required for all models that support quantization."""
 
-    hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None
+    hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None
     packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
     quant_config: QuantizationConfig | None = None
 
@@ -1054,8 +1053,8 @@ class SupportsRealtime(Protocol):
         cls,
         audio_stream: AsyncGenerator[np.ndarray, None],
         input_stream: asyncio.Queue[list[int]],
-        model_config: ModelConfig,
-    ) -> AsyncGenerator[PromptType, None]: ...
+        model_config: "ModelConfig",
+    ) -> AsyncGenerator["PromptType", None]: ...
 
 
 @overload
@@ -1124,8 +1123,8 @@ class SupportsTranscription(Protocol):
     @classmethod
     def get_generation_prompt(
         cls,
-        stt_params: SpeechToTextParams,
-    ) -> PromptType:
+        stt_params: "SpeechToTextParams",
+    ) -> "PromptType":
         """Get the prompt for the ASR model.
         The model has control over the construction, as long as it
         returns a valid PromptType."""
@@ -1163,8 +1162,8 @@ class SupportsTranscription(Protocol):
 
     @classmethod
     def get_speech_to_text_config(
-        cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
-    ) -> SpeechToTextConfig:
+        cls, model_config: "ModelConfig", task_type: Literal["transcribe", "translate"]
+    ) -> "SpeechToTextConfig":
         """Get the speech to text config for the ASR model."""
         ...
 
@@ -1172,8 +1171,8 @@ class SupportsTranscription(Protocol):
     def get_num_audio_tokens(
         cls,
         audio_duration_s: float,
-        stt_config: SpeechToTextConfig,
-        model_config: ModelConfig,
+        stt_config: "SpeechToTextConfig",
+        model_config: "ModelConfig",
     ) -> int | None:
         """
         Map from audio duration to number of audio tokens produced by the ASR
@@ -1202,8 +1201,8 @@ class SupportsTranscription(Protocol):
     def get_language_detection_prompt(
         cls,
         audio: np.ndarray,
-        stt_config: SpeechToTextConfig,
-    ) -> PromptType:
+        stt_config: "SpeechToTextConfig",
+    ) -> "PromptType":
         """Return a prompt that triggers language detection.
 
         Only needs to be implemented when

vllm/model_executor/models/interns1_pro.py

@@ -513,8 +513,6 @@ class InternS1ProMoeMixtureOfExperts(MixtureOfExperts):
                 moe.experts.update_expert_map()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.moe_layers = []
         example_moe = None
         for layer in self.language_model.model.layers:

vllm/model_executor/models/lfm2_moe.py

@@ -692,7 +692,6 @@ class Lfm2MoeForCausalLM(
         )
 
         # Set MoE hyperparameters
-        self.expert_weights = []
 
         self.moe_layers = []
         example_layer = None

vllm/model_executor/models/llama4.py

@@ -729,8 +729,6 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
         self.set_moe_parameters()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.moe_layers = []
         example_moe = None
         for layer in self.model.layers:

vllm/model_executor/models/mellum.py

@@ -227,8 +227,6 @@ class MellumForCausalLM(Qwen3MoeForCausalLM):
             self.model.make_empty_intermediate_tensors
         )
 
-        self.expert_weights = []
-
         self.moe_layers = []
         example_layer = None
         for layer in self.model.layers:

vllm/model_executor/models/mixtral.py

@@ -512,7 +512,6 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
             self.model.make_empty_intermediate_tensors
         )
 
-        self.expert_weights = []
         self.moe_layers = []
         example_moe = None
 

vllm/model_executor/models/nemotron_h.py

@@ -886,7 +886,6 @@ class NemotronHForCausalLM(
 
         # Set MoE hyperparameters
         if self.model.has_moe:
-            self.expert_weights = []
             self.num_expert_groups = config.n_group
 
             self.moe_layers = []

vllm/model_executor/models/openpangu.py

@@ -1289,7 +1289,6 @@ class OpenPanguMoEModel(OpenPanguModelBase, MixtureOfExperts):
         config = vllm_config.model_config.hf_config
 
         # Set MoE hyperparameters
-        self.expert_weights = []
         self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
         self.num_expert_groups = 1
 

vllm/model_executor/models/param2moe.py

@@ -751,7 +751,7 @@ class Param2MoEMixtureOfExperts(MixtureOfExperts):
         logical_to_physical_map: torch.Tensor,
         logical_replica_count: torch.Tensor,
     ) -> None:
-        self.expert_weights.clear()
+        self.expert_weights = []
         for layer_idx, layer in enumerate(self.moe_layers):
             if hasattr(layer, "get_expert_weights"):
                 self.expert_weights.append(layer.get_expert_weights())
@@ -832,7 +832,6 @@ class Param2MoEForCausalLM(
             self.model.make_empty_intermediate_tensors
         )
 
-        self.expert_weights: list[torch.Tensor] = []
         self.num_moe_layers: int = 0
         self.moe_layers: list = []
         self.moe_mlp_layers: list = []

vllm/model_executor/models/qwen3_5.py

@@ -765,8 +765,6 @@ class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
                 moe.experts.update_expert_map()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.moe_layers = []
         example_moe = None
         for layer in self.language_model.model.layers:

vllm/model_executor/models/qwen3_moe.py

@@ -703,7 +703,6 @@ class Qwen3MoeForCausalLM(
         )
 
         # Set MoE hyperparameters
-        self.expert_weights = []
 
         self.moe_layers = []
         example_layer = None

vllm/model_executor/models/qwen3_next.py

@@ -763,8 +763,6 @@ class QwenNextMixtureOfExperts(MixtureOfExperts):
                 moe.experts.update_expert_map()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.moe_layers = []
         example_moe = None
         for layer in self.model.layers:

vllm/model_executor/models/qwen3_vl_moe.py

@@ -373,8 +373,6 @@ class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts):
                 moe.experts.update_expert_map()
 
     def set_moe_parameters(self):
-        self.expert_weights = []
-
         self.moe_layers = []
         example_moe = None
         for layer in self.language_model.model.layers:

vllm/model_executor/models/step3p5.py

@@ -919,17 +919,17 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
         )
 
         # Set MoE hyperparameters
-        self.moe_layers: list[FusedMoEBlock] = []
+        self.moe_layers: list[MoERunner] = []
+        example_layer: FusedMoEBlock | None = None
         for layer in self.model.layers:
             if isinstance(layer, PPMissingLayer):
                 continue
             assert isinstance(layer, Step3p5DecoderLayer)
             if hasattr(layer, "moe") and isinstance(layer.moe, FusedMoEBlock):
-                self.moe_layers.append(layer.moe)
+                example_layer = layer.moe
+                self.moe_layers.append(layer.moe.experts)
 
-        self.expert_weights = []
         assert len(self.moe_layers) > 0, "No MoE layers found in the model."
-        example_layer = self.moe_layers[0]
         self.num_moe_layers = len(self.moe_layers)
         self.num_expert_groups = 1
         self.num_shared_experts = 0
@@ -959,24 +959,6 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_tokens(input_ids)
 
-    def set_eplb_state(
-        self,
-        expert_load_view: torch.Tensor,
-        logical_to_physical_map: torch.Tensor,
-        logical_replica_count: torch.Tensor,
-    ) -> None:
-        for layer_idx, layer in enumerate(self.moe_layers):
-            experts = layer.experts
-            assert isinstance(experts, MoERunner)
-            # Register the expert weights.
-            self.expert_weights.append(experts.get_expert_weights())
-            experts.set_eplb_state(
-                moe_layer_idx=layer_idx,
-                expert_load_view=expert_load_view,
-                logical_to_physical_map=logical_to_physical_map,
-                logical_replica_count=logical_replica_count,
-            )
-
     def update_physical_experts_metadata(
         self,
         num_physical_experts: int,

vllm/model_executor/models/transformers/moe.py

@@ -28,14 +28,9 @@ from vllm.config.utils import getattr_iter
 from vllm.distributed import get_dp_group, get_ep_group
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.model_executor.custom_op import PluggableLayer
-from vllm.model_executor.layers.fused_moe import (
-    FusedMoE,
-    MoERunner,
-    fused_moe_make_expert_params_mapping,
-)
+from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts
 from vllm.model_executor.models.interfaces import MixtureOfExperts
 from vllm.model_executor.models.utils import maybe_prefix
-from vllm.platforms import current_platform
 from vllm.utils.torch_utils import direct_register_custom_op
 
 from .utils import log_replacement
@@ -52,7 +47,7 @@ class TransformersMoEState:
 
 # --8<-- [start:transformers_fused_moe]
 @PluggableLayer.register("transformers_fused_moe")
-class TransformersFusedMoE(MoERunner):
+class TransformersMoERunner(MoERunner):
     """Custom FusedMoE for the Transformers modeling backend."""
 
     # --8<-- [end:transformers_fused_moe]
@@ -93,7 +88,7 @@ class TransformersFusedMoE(MoERunner):
         return self.routed_experts.load_weights(weights)
 
 
-def transformers_moe_forward(
+def _transformers_moe_forward(
     hidden_states: torch.Tensor,
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
@@ -106,7 +101,7 @@ def transformers_moe_forward(
     return self._forward_super(hidden_states, topk_weights)
 
 
-def transformers_moe_forward_fake(
+def _transformers_moe_forward_fake(
     hidden_states: torch.Tensor,
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
@@ -117,10 +112,9 @@ def transformers_moe_forward_fake(
 
 direct_register_custom_op(
     op_name="transformers_moe_forward",
-    op_func=transformers_moe_forward,
+    op_func=_transformers_moe_forward,
     mutates_args=["hidden_states"],
-    fake_impl=transformers_moe_forward_fake,
-    dispatch_key=current_platform.dispatch_key,
+    fake_impl=_transformers_moe_forward_fake,
     tags=(torch.Tag.needs_fixed_stride_order,),
 )
 
@@ -131,20 +125,6 @@ class MoEMixin(MixtureOfExperts):
         # Skip MixtureOfExperts.__init__ and call the next class in MRO
         super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix)
 
-    def set_eplb_state(
-        self,
-        expert_load_view: torch.Tensor,
-        logical_to_physical_map: torch.Tensor,
-        logical_replica_count: torch.Tensor,
-    ):
-        for moe_layer_idx, mlp_layer in enumerate(self.mlp_moe_layers):
-            mlp_layer.experts.set_eplb_state(
-                moe_layer_idx=moe_layer_idx,
-                expert_load_view=expert_load_view,
-                logical_to_physical_map=logical_to_physical_map,
-                logical_replica_count=logical_replica_count,
-            )
-
     def update_physical_experts_metadata(
         self,
         num_physical_experts: int,
@@ -154,7 +134,7 @@ class MoEMixin(MixtureOfExperts):
         self.num_physical_experts = num_physical_experts
         self.num_local_physical_experts = num_local_physical_experts
         self.num_redundant_experts = num_physical_experts - self.num_logical_experts
-        for mlp in self.mlp_moe_layers:
+        for mlp in self.mlp_layers:
             mlp.n_local_physical_experts = num_local_physical_experts
             mlp.n_physical_experts = num_physical_experts
             mlp.n_redundant_experts = self.num_redundant_experts
@@ -185,7 +165,7 @@ class MoEMixin(MixtureOfExperts):
         num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
         for gate_proj, down_proj, up_proj in ckpt_names:
             expert_mapping.extend(
-                fused_moe_make_expert_params_mapping(
+                RoutedExperts.make_expert_params_mapping(
                     self,
                     ckpt_gate_proj_name=gate_proj,
                     ckpt_down_proj_name=down_proj,
@@ -248,10 +228,8 @@ class MoEMixin(MixtureOfExperts):
         # MixtureOfExperts mixin settings
         ep_size = get_ep_group().world_size
 
-        self.mlp_moe_layers = []  # Used for MixtureOfExperts methods
+        self.mlp_layers = []  # Used for MixtureOfExperts methods
         self.moe_layers = []
-        self.expert_weights = []
-        self.num_moe_layers = 0
         self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
         self.num_logical_experts = num_experts
         self.num_physical_experts = num_experts + num_redundant_experts
@@ -335,19 +313,18 @@ class MoEMixin(MixtureOfExperts):
                             custom_routing_function,
                             moe_state=moe_state,
                         ),
-                        runner_cls=TransformersFusedMoE,
+                        runner_cls=TransformersMoERunner,
                         runner_args={"moe_state": moe_state},
                     )
                     mlp.experts = fused_experts
                     log_replacement(qual_name, experts, fused_experts)
                     # Update MixtureOfExperts mixin state
-                    self.mlp_moe_layers.append(mlp)
+                    self.mlp_layers.append(mlp)
                     self.moe_layers.append(fused_experts)
-                    self.expert_weights.append(fused_experts.get_expert_weights())
-                    self.num_moe_layers += 1
                 else:
                     _recursive_replace(child_module, prefix=qual_name)
 
         _recursive_replace(self.model, prefix="model")
+        self.num_moe_layers = len(self.moe_layers)
         # Continue with the replacement of layers in Base
         super().recursive_replace()