diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index d526f57d3d..a465c6b563 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/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 = [] diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 1564ee733f..369b7c3b3a 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/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) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 88f33ac021..f73d9f9c3e 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/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 diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 8d20e0b5c6..814118f8a7 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/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 = [] diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index e1b9ca9bf5..c2d9f92a66 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/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 diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 03e67c4ada..9cf86a5ba8 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/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 diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index bad7e061cc..30c379d86c 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/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 diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 98cc9a50ad..8226b65c45 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/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 diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index b4d0fe9668..432fa5e6fa 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/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 = [] diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index 4813af5f03..222705c14e 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/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 diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index d87ad26828..b255b67d88 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/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 diff --git a/vllm/model_executor/models/glm_ocr_mtp.py b/vllm/model_executor/models/glm_ocr_mtp.py index 3d283c101c..9b2369f93d 100644 --- a/vllm/model_executor/models/glm_ocr_mtp.py +++ b/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) diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index ec3cfbd017..4f70a96628 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/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 diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 7653cddd6c..a4b52e20bd 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/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 diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index ad3d01ae9e..29603318c1 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/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 diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 36f669179c..c04b472945 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/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: diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 9ca7fb7aaa..94f7f4e289 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/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 diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 9222405ba6..71df54a424 100644 --- a/vllm/model_executor/models/llama4.py +++ b/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: diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py index bdbf0df7fd..c20fa00e3d 100644 --- a/vllm/model_executor/models/mellum.py +++ b/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: diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 53c1c87cfc..57eb820ad9 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/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 diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 769504c0d0..bd5cd358d8 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/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 = [] diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 8432566a15..91120840bd 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/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 diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index e8ea2dbc0e..3386f9545f 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/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 = [] diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index b00b495868..480ef3678c 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/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: diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 6980184cc8..b7a78acc0e 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/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 diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 2c7667a416..74c2b1e44a 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/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: diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 298863209d..5291874dd5 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/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: diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 7a60946ba5..f8bd529e27 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/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, diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 372c5b1ec1..3e04aaa074 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/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()