@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d8eb734d94

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

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

Harry Mellor committed Jun 27, 2026 at 00:16 UTC d8eb734d94fea27cfcc95a22f3cc2a249e0996c7

29 files changed +80 -156

vllm/model_executor/models/AXK1.py

-2

@@ -1088,8 +1088,6 @@ class AXK1ForCausalLM(

1088 self.set_moe_parameters()

1089

1090 def set_moe_parameters(self):

1091 - self.expert_weights = []

1092 -

1091 self.num_expert_groups = getattr(self.config, "n_group", 1)

1092

1093 self.moe_layers = []

vllm/model_executor/models/afmoe.py

+20 -16

@@ -42,6 +42,7 @@ from vllm.model_executor.model_loader.weight_utils import (

42 )

43 from vllm.model_executor.models.interfaces import (

44 EagleModelMixin,

45 + MixtureOfExperts,

46 SupportsEagle3,

47 SupportsLoRA,

48 SupportsPP,

@@ -595,7 +596,9 @@ class AfmoeModel(nn.Module, EagleModelMixin):

596 return loaded_params

597

598

598 -class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

599 +class AfmoeForCausalLM(

600 + nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA, MixtureOfExperts

601 +):

602 packed_modules_mapping = {

603 "qkv_proj": [

604 "q_proj",

@@ -635,8 +638,6 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

638 self.make_empty_intermediate_tensors = (

639 self.model.make_empty_intermediate_tensors

640 )

638 - self.expert_weights = []

639 -

641 # Set MoE hyperparameters

642 self.num_moe_layers = config.num_hidden_layers - config.num_dense_layers

643 self.num_expert_groups = config.n_group

@@ -663,21 +664,24 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

664 self.num_shared_experts = example_moe.n_shared_experts

665 self.num_redundant_experts = example_moe.n_redundant_experts

666

666 - def set_eplb_state(

667 + def update_physical_experts_metadata(

668 self,

668 - expert_load_view: torch.Tensor,

669 - logical_to_physical_map: torch.Tensor,

670 - logical_replica_count: torch.Tensor,

669 + num_physical_experts: int,

670 + num_local_physical_experts: int,

671 ) -> None:

672 - for layer_idx, layer in enumerate(self.moe_layers):

673 - # Register the expert weights.

674 - self.expert_weights.append(layer.get_expert_weights())

675 - layer.set_eplb_state(

676 - moe_layer_idx=layer_idx,

677 - expert_load_view=expert_load_view,

678 - logical_to_physical_map=logical_to_physical_map,

679 - logical_replica_count=logical_replica_count,

680 - )

672 + assert self.num_local_physical_experts == num_local_physical_experts

673 + self.num_physical_experts = num_physical_experts

674 + self.num_local_physical_experts = num_local_physical_experts

675 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts

676 + for layer in self.model.layers:

677 + if isinstance(layer, PPMissingLayer):

678 + continue

679 + if layer.moe_enabled:

680 + moe = layer.mlp

681 + moe.n_local_physical_experts = num_local_physical_experts

682 + moe.n_physical_experts = num_physical_experts

683 + moe.n_redundant_experts = self.num_redundant_experts

684 + moe.experts.update_expert_map()

685

686 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

687 return self.model.embed_input_ids(input_ids)

vllm/model_executor/models/deepseek_mtp.py

-1

@@ -217,7 +217,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

217 self.set_moe_parameters()

218

219 def set_moe_parameters(self):

220 - self.expert_weights = []

220 self.num_moe_layers = self.config.num_nextn_predict_layers

221 self.num_expert_groups = self.config.n_group

222

vllm/model_executor/models/deepseek_v2.py

-2

@@ -1684,8 +1684,6 @@ class DeepseekV2ForCausalLM(

1684 self.set_moe_parameters()

1685

1686 def set_moe_parameters(self):

1687 - self.expert_weights = []

1688 -

1687 self.num_expert_groups = getattr(self.config, "n_group", 1)

1688

1689 self.moe_layers = []

vllm/model_executor/models/ernie45_moe.py

-2

@@ -656,8 +656,6 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe

656 self.model.make_empty_intermediate_tensors

657 )

658

659 - self.expert_weights = []

660 -

659 # Set MoE hyperparameters

660 moe_layers_indices = [

661 i

vllm/model_executor/models/gemma4.py

-1

@@ -1562,7 +1562,6 @@ class Gemma4ForCausalLM(

1562 )

1563

1564 # --- MixtureOfExperts protocol ---

1565 - self.expert_weights: list[list[torch.Tensor]] = []

1565 self.moe_layers: list[nn.Module] = []

1566 example_moe: Gemma4MoE | None = None

1567

vllm/model_executor/models/gemma4_mm.py

-1

@@ -1114,7 +1114,6 @@ class Gemma4ForConditionalGeneration(

1114 )

1115

1116 # --- MixtureOfExperts delegation to language_model ---

1117 - self.expert_weights = self.language_model.expert_weights

1117 self.moe_layers = self.language_model.moe_layers

1118 self.num_moe_layers = self.language_model.num_moe_layers

1119 self.num_logical_experts = self.language_model.num_logical_experts

vllm/model_executor/models/glm4_moe.py

-2

@@ -655,8 +655,6 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, Glm4MixtureOfExper

655 self.make_empty_intermediate_tensors = (

656 self.model.make_empty_intermediate_tensors

657 )

658 - self.expert_weights = []

659 -

658 # Set MoE hyperparameters

659 self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace

660 self.num_expert_groups = config.n_group

vllm/model_executor/models/glm4_moe_lite.py

-2

@@ -573,8 +573,6 @@ class Glm4MoeLiteForCausalLM(

573 self.set_moe_parameters()

574

575 def set_moe_parameters(self):

576 - self.expert_weights = []

577 -

576 self.num_expert_groups = getattr(self.config, "n_group", 1)

577

578 self.moe_layers = []

vllm/model_executor/models/glm4_moe_lite_mtp.py

-2

@@ -209,8 +209,6 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

209 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

210 )

211

212 - self.expert_weights = []

213 -

212 # Set MoE hyperparameters

213 self.num_moe_layers = self.config.num_nextn_predict_layers

214 self.num_expert_groups = self.config.n_group

vllm/model_executor/models/glm4_moe_mtp.py

-2

@@ -195,8 +195,6 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):

195 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

196 )

197

198 - self.expert_weights = []

199 -

198 # Set MoE hyperparameters

199 self.num_moe_layers = self.config.num_nextn_predict_layers

200 self.num_expert_groups = self.config.n_group

vllm/model_executor/models/glm_ocr_mtp.py

-1

@@ -134,7 +134,6 @@ class GlmOcrMTP(nn.Module, SupportsPP):

134 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

135 )

136

137 - self.expert_weights = []

137 self.num_layers = self.config.num_nextn_predict_layers

138 for layer in self.model.layers.values():

139 assert isinstance(layer, GlmOcrMultiTokenPredictorLayer)

vllm/model_executor/models/hunyuan_v1.py

-1

@@ -991,7 +991,6 @@ class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts):

991 super().__init__(vllm_config=vllm_config, prefix=prefix)

992

993 # Set MoE hyperparameters

994 - self.expert_weights = []

994 self.num_expert_groups = 1

995 self.moe_layers = []

996 example_layer = None

vllm/model_executor/models/hy_v3.py

+2 -3

@@ -68,7 +68,7 @@ from vllm.model_executor.model_loader.weight_utils import (

68 from vllm.sequence import IntermediateTensors

69 from vllm.transformers_utils.configs.hy_v3 import HYV3Config

71 -from .interfaces import SupportsLoRA, SupportsPP

71 +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP

72 from .utils import (

73 AutoWeightsLoader,

74 PPMissingLayer,

@@ -392,7 +392,7 @@ class HYV3DecoderLayer(nn.Module):

392

393

394 @support_torch_compile

395 -class HYV3Model(nn.Module):

395 +class HYV3Model(nn.Module, MixtureOfExperts):

396 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

397 super().__init__()

398

@@ -429,7 +429,6 @@ class HYV3Model(nn.Module):

429 )

430

431 # Set MoE hyperparameters

432 - self.expert_weights = []

432 self.num_expert_groups = 1

433 self.moe_layers = []

434 example_layer = None

vllm/model_executor/models/interfaces.py

+41 -42

@@ -29,36 +29,34 @@ from torch import Tensor

29 from transformers.models.whisper.tokenization_whisper import LANGUAGES

30 from typing_extensions import Self, TypeIs

32 -from vllm.config import ModelConfig, SpeechToTextConfig, SpeechToTextParams

33 -from vllm.inputs import PromptType, TokensPrompt

32 from vllm.logger import init_logger

35 -from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc

33 from vllm.model_executor.layers.quantization import QuantizationConfig

37 -from vllm.tasks import ScoreType

34 from vllm.utils.collection_utils import common_prefix

35 from vllm.utils.func_utils import supports_kw

41 -from .interfaces_base import VllmModel

42 -

37 if TYPE_CHECKING:

44 - from vllm.config import VllmConfig

38 + from vllm.config import (

39 + ModelConfig,

40 + SpeechToTextConfig,

41 + SpeechToTextParams,

42 + VllmConfig,

43 + )

44 + from vllm.inputs import PromptType, TokensPrompt

45 from vllm.lora.model_manager import LoRAModelManager

46 + from vllm.model_executor.layers.fused_moe import MoERunner

47 + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc

48 + from vllm.model_executor.models.interfaces_base import VllmModel

49 from vllm.model_executor.models.utils import WeightsMapper

50 from vllm.multimodal.inputs import MultiModalFeatureSpec

51 from vllm.multimodal.registry import _ProcessorFactories

52 from vllm.sequence import IntermediateTensors

53 + from vllm.tasks import ScoreType

54 from vllm.v1.worker.encoder_cudagraph_defs import (

55 EncoderCudaGraphCaptureInputs,

56 EncoderCudaGraphConfig,

57 EncoderCudaGraphReplayBuffers,

58 EncoderItemSpec,

59 )

56 -else:

57 - VllmConfig = object

58 - WeightsMapper = object

59 - MultiModalFeatureSpec = object

60 - _ProcessorFactories = object

61 - IntermediateTensors = object

61 logger = init_logger(__name__)

@@ -89,7 +87,7 @@ def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:

89 # Cache results of `SupportsMultiModal.get_language_model`

92 -_language_model_by_module = dict[nn.Module, VllmModel]()

90 +_language_model_by_module = dict[nn.Module, "VllmModel"]()

93 @runtime_checkable

@@ -123,7 +121,7 @@ class SupportsMultiModal(Protocol):

121 in their raw form and not input embeddings.

122 """

123

126 - _processor_factory: ClassVar[_ProcessorFactories]

124 + _processor_factory: ClassVar["_ProcessorFactories"]

125 """

126 Set internally by `MultiModalRegistry.register_processor`.

127 """

@@ -175,7 +173,7 @@ class SupportsMultiModal(Protocol):

173 self._has_oov_mm_tokens,

174 )

175

178 - def get_language_model(self) -> VllmModel:

176 + def get_language_model(self) -> "VllmModel":

177 """

178 Returns the underlying language model used for text generation.

179

@@ -216,7 +214,7 @@ class SupportsMultiModal(Protocol):

214 @contextmanager

215 def _mark_language_model(

216 self,

219 - vllm_config: VllmConfig,

217 + vllm_config: "VllmConfig",

218 *,

219 targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,

220 ):

@@ -251,7 +249,7 @@ class SupportsMultiModal(Protocol):

249 @contextmanager

250 def _mark_tower_model(

251 self,

254 - vllm_config: VllmConfig,

252 + vllm_config: "VllmConfig",

253 modalities: set[str] | str,

254 *,

255 targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,

@@ -295,7 +293,7 @@ class SupportsMultiModal(Protocol):

293 @contextmanager

294 def _mark_composite_model(

295 self,

298 - vllm_config: VllmConfig,

296 + vllm_config: "VllmConfig",

297 *,

298 language_targets: type[nn.Module] | tuple[type[nn.Module], ...],

299 tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]],

@@ -513,7 +511,7 @@ class SupportsScoreTemplate(Protocol):

511 ...

512

513 @classmethod

516 - def post_process_tokens(cls, prompt: TokensPrompt) -> None:

514 + def post_process_tokens(cls, prompt: "TokensPrompt") -> None:

515 """

516 Perform architecture-specific manipulations on the input tokens.

517 """

@@ -633,7 +631,7 @@ class SupportsPP(Protocol):

631 batch_size: int,

632 dtype: torch.dtype,

633 device: torch.device,

636 - ) -> IntermediateTensors:

634 + ) -> "IntermediateTensors":

635 """Called when PP rank > 0 for profiling purposes."""

636 ...

637

@@ -642,8 +640,8 @@ class SupportsPP(Protocol):

640 input_ids: Tensor | None,

641 positions: Tensor,

642 *,

645 - intermediate_tensors: IntermediateTensors | None,

646 - ) -> IntermediateTensors | None:

643 + intermediate_tensors: "IntermediateTensors | None",

644 + ) -> "IntermediateTensors | None":

645 """

646 Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when

647 PP rank > 0.

@@ -665,15 +663,15 @@ class _SupportsPPType(Protocol):

663 batch_size: int,

664 dtype: torch.dtype,

665 device: torch.device,

668 - ) -> IntermediateTensors: ...

666 + ) -> "IntermediateTensors": ...

667

668 def forward(

669 self,

670 input_ids: Tensor | None,

671 positions: Tensor,

672 *,

675 - intermediate_tensors: IntermediateTensors | None,

676 - ) -> Tensor | IntermediateTensors: ...

673 + intermediate_tensors: "IntermediateTensors | None",

674 + ) -> "Tensor | IntermediateTensors": ...

675

676

677 @overload

@@ -803,7 +801,7 @@ class IsHybrid(Protocol):

801 @classmethod

802 def get_mamba_state_shape_from_config(

803 cls,

806 - vllm_config: VllmConfig,

804 + vllm_config: "VllmConfig",

805 ) -> tuple[tuple[int, int], tuple[int, int, int]]:

806 """Calculate shapes for Mamba's convolutional and state caches.

807

@@ -818,7 +816,7 @@ class IsHybrid(Protocol):

816 ...

817

818 @classmethod

821 - def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:

819 + def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]:

820 """Calculate copy-function callables for each Mamba state.

821

822 Returns:

@@ -883,7 +881,7 @@ class MixtureOfExperts(Protocol):

881 num_redundant_experts: int

882 """Number of redundant experts in this model."""

883

886 - moe_layers: Iterable[nn.Module]

884 + moe_layers: Iterable["MoERunner"]

885 """List of MoE layers in this model."""

886

887 def set_eplb_state(

@@ -908,6 +906,7 @@ class MixtureOfExperts(Protocol):

906 logical_to_physical_map: Mapping from logical to physical experts.

907 logical_replica_count: Count of replicas for each logical expert.

908 """

909 + self.expert_weights = []

910 for layer_idx, layer in enumerate(self.moe_layers):

911 # Register the expert weights.

912 self.expert_weights.append(layer.get_expert_weights())

@@ -982,7 +981,7 @@ def supports_mamba_prefix_caching(

981 class SupportsCrossEncoding(Protocol):

982 """The interface required for all models that support cross encoding."""

983

985 - score_type: ClassVar[ScoreType] = "cross-encoder"

984 + score_type: ClassVar["ScoreType"] = "cross-encoder"

985

986

987 @runtime_checkable

@@ -994,13 +993,13 @@ class SupportsLateInteraction(Protocol):

993 MaxSim (max over document tokens, sum over query tokens).

994 """

995

997 - score_type: ClassVar[ScoreType] = "late-interaction"

996 + score_type: ClassVar["ScoreType"] = "late-interaction"

997

998

999 class SupportsQuant:

1000 """The interface required for all models that support quantization."""

1001

1003 - hf_to_vllm_mapper: ClassVar[WeightsMapper | None] = None

1002 + hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None

1003 packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None

1004 quant_config: QuantizationConfig | None = None

1005

@@ -1054,8 +1053,8 @@ class SupportsRealtime(Protocol):

1053 cls,

1054 audio_stream: AsyncGenerator[np.ndarray, None],

1055 input_stream: asyncio.Queue[list[int]],

1057 - model_config: ModelConfig,

1058 - ) -> AsyncGenerator[PromptType, None]: ...

1056 + model_config: "ModelConfig",

1057 + ) -> AsyncGenerator["PromptType", None]: ...

1058

1059

1060 @overload

@@ -1124,8 +1123,8 @@ class SupportsTranscription(Protocol):

1123 @classmethod

1124 def get_generation_prompt(

1125 cls,

1127 - stt_params: SpeechToTextParams,

1128 - ) -> PromptType:

1126 + stt_params: "SpeechToTextParams",

1127 + ) -> "PromptType":

1128 """Get the prompt for the ASR model.

1129 The model has control over the construction, as long as it

1130 returns a valid PromptType."""

@@ -1163,8 +1162,8 @@ class SupportsTranscription(Protocol):

1162

1163 @classmethod

1164 def get_speech_to_text_config(

1166 - cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]

1167 - ) -> SpeechToTextConfig:

1165 + cls, model_config: "ModelConfig", task_type: Literal["transcribe", "translate"]

1166 + ) -> "SpeechToTextConfig":

1167 """Get the speech to text config for the ASR model."""

1168 ...

1169

@@ -1172,8 +1171,8 @@ class SupportsTranscription(Protocol):

1171 def get_num_audio_tokens(

1172 cls,

1173 audio_duration_s: float,

1175 - stt_config: SpeechToTextConfig,

1176 - model_config: ModelConfig,

1174 + stt_config: "SpeechToTextConfig",

1175 + model_config: "ModelConfig",

1176 ) -> int | None:

1177 """

1178 Map from audio duration to number of audio tokens produced by the ASR

@@ -1202,8 +1201,8 @@ class SupportsTranscription(Protocol):

1201 def get_language_detection_prompt(

1202 cls,

1203 audio: np.ndarray,

1205 - stt_config: SpeechToTextConfig,

1206 - ) -> PromptType:

1204 + stt_config: "SpeechToTextConfig",

1205 + ) -> "PromptType":

1206 """Return a prompt that triggers language detection.

1207

1208 Only needs to be implemented when

vllm/model_executor/models/interns1_pro.py

-2

@@ -513,8 +513,6 @@ class InternS1ProMoeMixtureOfExperts(MixtureOfExperts):

513 moe.experts.update_expert_map()

514

515 def set_moe_parameters(self):

516 - self.expert_weights = []

517 -

516 self.moe_layers = []

517 example_moe = None

518 for layer in self.language_model.model.layers:

vllm/model_executor/models/lfm2_moe.py

-1

@@ -692,7 +692,6 @@ class Lfm2MoeForCausalLM(

692 )

693

694 # Set MoE hyperparameters

695 - self.expert_weights = []

695

696 self.moe_layers = []

697 example_layer = None

vllm/model_executor/models/llama4.py

-2

@@ -729,8 +729,6 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):

729 self.set_moe_parameters()

730

731 def set_moe_parameters(self):

732 - self.expert_weights = []

733 -

732 self.moe_layers = []

733 example_moe = None

734 for layer in self.model.layers:

vllm/model_executor/models/mellum.py

-2

@@ -227,8 +227,6 @@ class MellumForCausalLM(Qwen3MoeForCausalLM):

227 self.model.make_empty_intermediate_tensors

228 )

229

230 - self.expert_weights = []

231 -

230 self.moe_layers = []

231 example_layer = None

232 for layer in self.model.layers:

vllm/model_executor/models/mixtral.py

-1

@@ -512,7 +512,6 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):

512 self.model.make_empty_intermediate_tensors

513 )

514

515 - self.expert_weights = []

515 self.moe_layers = []

516 example_moe = None

517

vllm/model_executor/models/nemotron_h.py

-1

@@ -886,7 +886,6 @@ class NemotronHForCausalLM(

886

887 # Set MoE hyperparameters

888 if self.model.has_moe:

889 - self.expert_weights = []

889 self.num_expert_groups = config.n_group

890

891 self.moe_layers = []

vllm/model_executor/models/openpangu.py

-1

@@ -1289,7 +1289,6 @@ class OpenPanguMoEModel(OpenPanguModelBase, MixtureOfExperts):

1289 config = vllm_config.model_config.hf_config

1290

1291 # Set MoE hyperparameters

1292 - self.expert_weights = []

1292 self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace

1293 self.num_expert_groups = 1

1294

vllm/model_executor/models/param2moe.py

+1 -2

@@ -751,7 +751,7 @@ class Param2MoEMixtureOfExperts(MixtureOfExperts):

751 logical_to_physical_map: torch.Tensor,

752 logical_replica_count: torch.Tensor,

753 ) -> None:

754 - self.expert_weights.clear()

754 + self.expert_weights = []

755 for layer_idx, layer in enumerate(self.moe_layers):

756 if hasattr(layer, "get_expert_weights"):

757 self.expert_weights.append(layer.get_expert_weights())

@@ -832,7 +832,6 @@ class Param2MoEForCausalLM(

832 self.model.make_empty_intermediate_tensors

833 )

834

835 - self.expert_weights: list[torch.Tensor] = []

835 self.num_moe_layers: int = 0

836 self.moe_layers: list = []

837 self.moe_mlp_layers: list = []

vllm/model_executor/models/qwen3_5.py

-2

@@ -765,8 +765,6 @@ class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):

765 moe.experts.update_expert_map()

766

767 def set_moe_parameters(self):

768 - self.expert_weights = []

769 -

768 self.moe_layers = []

769 example_moe = None

770 for layer in self.language_model.model.layers:

vllm/model_executor/models/qwen3_moe.py

-1

@@ -703,7 +703,6 @@ class Qwen3MoeForCausalLM(

703 )

704

705 # Set MoE hyperparameters

706 - self.expert_weights = []

706

707 self.moe_layers = []

708 example_layer = None

vllm/model_executor/models/qwen3_next.py

-2

@@ -763,8 +763,6 @@ class QwenNextMixtureOfExperts(MixtureOfExperts):

763 moe.experts.update_expert_map()

764

765 def set_moe_parameters(self):

766 - self.expert_weights = []

767 -

766 self.moe_layers = []

767 example_moe = None

768 for layer in self.model.layers:

vllm/model_executor/models/qwen3_vl_moe.py

-2

@@ -373,8 +373,6 @@ class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts):

373 moe.experts.update_expert_map()

374

375 def set_moe_parameters(self):

376 - self.expert_weights = []

377 -

376 self.moe_layers = []

377 example_moe = None

378 for layer in self.language_model.model.layers:

vllm/model_executor/models/step3p5.py

+4 -22

@@ -919,17 +919,17 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

919 )

920

921 # Set MoE hyperparameters

922 - self.moe_layers: list[FusedMoEBlock] = []

922 + self.moe_layers: list[MoERunner] = []

923 + example_layer: FusedMoEBlock | None = None

924 for layer in self.model.layers:

925 if isinstance(layer, PPMissingLayer):

926 continue

927 assert isinstance(layer, Step3p5DecoderLayer)

928 if hasattr(layer, "moe") and isinstance(layer.moe, FusedMoEBlock):

928 - self.moe_layers.append(layer.moe)

929 + example_layer = layer.moe

930 + self.moe_layers.append(layer.moe.experts)

931

930 - self.expert_weights = []

932 assert len(self.moe_layers) > 0, "No MoE layers found in the model."

932 - example_layer = self.moe_layers[0]

933 self.num_moe_layers = len(self.moe_layers)

934 self.num_expert_groups = 1

935 self.num_shared_experts = 0

@@ -959,24 +959,6 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

959 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

960 return self.model.embed_tokens(input_ids)

961

962 - def set_eplb_state(

963 - self,

964 - expert_load_view: torch.Tensor,

965 - logical_to_physical_map: torch.Tensor,

966 - logical_replica_count: torch.Tensor,

967 - ) -> None:

968 - for layer_idx, layer in enumerate(self.moe_layers):

969 - experts = layer.experts

970 - assert isinstance(experts, MoERunner)

971 - # Register the expert weights.

972 - self.expert_weights.append(experts.get_expert_weights())

973 - experts.set_eplb_state(

974 - moe_layer_idx=layer_idx,

975 - expert_load_view=expert_load_view,

976 - logical_to_physical_map=logical_to_physical_map,

977 - logical_replica_count=logical_replica_count,

978 - )

979 -

962 def update_physical_experts_metadata(

963 self,

964 num_physical_experts: int,

vllm/model_executor/models/transformers/moe.py

+12 -35

@@ -28,14 +28,9 @@ from vllm.config.utils import getattr_iter

28 from vllm.distributed import get_dp_group, get_ep_group

29 from vllm.forward_context import ForwardContext, get_forward_context

30 from vllm.model_executor.custom_op import PluggableLayer

31 -from vllm.model_executor.layers.fused_moe import (

32 - FusedMoE,

33 - MoERunner,

34 - fused_moe_make_expert_params_mapping,

35 -)

31 +from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts

32 from vllm.model_executor.models.interfaces import MixtureOfExperts

33 from vllm.model_executor.models.utils import maybe_prefix

38 -from vllm.platforms import current_platform

34 from vllm.utils.torch_utils import direct_register_custom_op

36 from .utils import log_replacement

@@ -52,7 +47,7 @@ class TransformersMoEState:

48 # --8<-- [start:transformers_fused_moe]

49 @PluggableLayer.register("transformers_fused_moe")

55 -class TransformersFusedMoE(MoERunner):

50 +class TransformersMoERunner(MoERunner):

51 """Custom FusedMoE for the Transformers modeling backend."""

53 # --8<-- [end:transformers_fused_moe]

@@ -93,7 +88,7 @@ class TransformersFusedMoE(MoERunner):

88 return self.routed_experts.load_weights(weights)

96 -def transformers_moe_forward(

91 +def _transformers_moe_forward(

92 hidden_states: torch.Tensor,

93 topk_ids: torch.Tensor,

94 topk_weights: torch.Tensor,

@@ -106,7 +101,7 @@ def transformers_moe_forward(

101 return self._forward_super(hidden_states, topk_weights)

102

103

109 -def transformers_moe_forward_fake(

104 +def _transformers_moe_forward_fake(

105 hidden_states: torch.Tensor,

106 topk_ids: torch.Tensor,

107 topk_weights: torch.Tensor,

@@ -117,10 +112,9 @@ def transformers_moe_forward_fake(

112

113 direct_register_custom_op(

114 op_name="transformers_moe_forward",

120 - op_func=transformers_moe_forward,

115 + op_func=_transformers_moe_forward,

116 mutates_args=["hidden_states"],

122 - fake_impl=transformers_moe_forward_fake,

123 - dispatch_key=current_platform.dispatch_key,

117 + fake_impl=_transformers_moe_forward_fake,

118 tags=(torch.Tag.needs_fixed_stride_order,),

119 )

120

@@ -131,20 +125,6 @@ class MoEMixin(MixtureOfExperts):

125 # Skip MixtureOfExperts.__init__ and call the next class in MRO

126 super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix)

127

134 - def set_eplb_state(

135 - self,

136 - expert_load_view: torch.Tensor,

137 - logical_to_physical_map: torch.Tensor,

138 - logical_replica_count: torch.Tensor,

139 - ):

140 - for moe_layer_idx, mlp_layer in enumerate(self.mlp_moe_layers):

141 - mlp_layer.experts.set_eplb_state(

142 - moe_layer_idx=moe_layer_idx,

143 - expert_load_view=expert_load_view,

144 - logical_to_physical_map=logical_to_physical_map,

145 - logical_replica_count=logical_replica_count,

146 - )

147 -

128 def update_physical_experts_metadata(

129 self,

130 num_physical_experts: int,

@@ -154,7 +134,7 @@ class MoEMixin(MixtureOfExperts):

134 self.num_physical_experts = num_physical_experts

135 self.num_local_physical_experts = num_local_physical_experts

136 self.num_redundant_experts = num_physical_experts - self.num_logical_experts

157 - for mlp in self.mlp_moe_layers:

137 + for mlp in self.mlp_layers:

138 mlp.n_local_physical_experts = num_local_physical_experts

139 mlp.n_physical_experts = num_physical_experts

140 mlp.n_redundant_experts = self.num_redundant_experts

@@ -185,7 +165,7 @@ class MoEMixin(MixtureOfExperts):

165 num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts

166 for gate_proj, down_proj, up_proj in ckpt_names:

167 expert_mapping.extend(

188 - fused_moe_make_expert_params_mapping(

168 + RoutedExperts.make_expert_params_mapping(

169 self,

170 ckpt_gate_proj_name=gate_proj,

171 ckpt_down_proj_name=down_proj,

@@ -248,10 +228,8 @@ class MoEMixin(MixtureOfExperts):

228 # MixtureOfExperts mixin settings

229 ep_size = get_ep_group().world_size

230

251 - self.mlp_moe_layers = [] # Used for MixtureOfExperts methods

231 + self.mlp_layers = [] # Used for MixtureOfExperts methods

232 self.moe_layers = []

253 - self.expert_weights = []

254 - self.num_moe_layers = 0

233 self.num_expert_groups = 1 if num_expert_group is None else num_expert_group

234 self.num_logical_experts = num_experts

235 self.num_physical_experts = num_experts + num_redundant_experts

@@ -335,19 +313,18 @@ class MoEMixin(MixtureOfExperts):

313 custom_routing_function,

314 moe_state=moe_state,

315 ),

338 - runner_cls=TransformersFusedMoE,

316 + runner_cls=TransformersMoERunner,

317 runner_args={"moe_state": moe_state},

318 )

319 mlp.experts = fused_experts

320 log_replacement(qual_name, experts, fused_experts)

321 # Update MixtureOfExperts mixin state

344 - self.mlp_moe_layers.append(mlp)

322 + self.mlp_layers.append(mlp)

323 self.moe_layers.append(fused_experts)

346 - self.expert_weights.append(fused_experts.get_expert_weights())

347 - self.num_moe_layers += 1

324 else:

325 _recursive_replace(child_module, prefix=qual_name)

326

327 _recursive_replace(self.model, prefix="model")

328 + self.num_moe_layers = len(self.moe_layers)

329 # Continue with the replacement of layers in Base

330 super().recursive_replace()