diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index 4ffb081ca3..8f04ce3577 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -10,11 +10,24 @@ import torch.nn as nn from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import is_mixture_of_experts +from vllm.model_executor.models.interfaces import ( + SupportsMultiModal, + is_mixture_of_experts, +) logger = init_logger(__name__) +def _unwrap_moe(model: nn.Module) -> nn.Module: + # VLM wrappers (e.g. KimiK25ForConditionalGeneration) hold the MoE + # language model under `.language_model` but don't implement + # MixtureOfExperts themselves. Mirror the V1 path + # (see vllm/v1/worker/gpu_model_runner.py, PR #39805). + if not is_mixture_of_experts(model) and isinstance(model, SupportsMultiModal): + return model.get_language_model() + return model + + def step_eplb_after(*, is_dummy: bool = False) -> Callable: """Step EPLB after a model runner method completes successfully.""" @@ -89,6 +102,7 @@ class EPLBController: if not self.parallel_config.enable_eplb or load_dummy_weights: return False + model = _unwrap_moe(model) if not is_mixture_of_experts(model): return False @@ -128,6 +142,7 @@ class EPLBController: expanded_physical_to_logical: torch.Tensor, old_num_physical_experts: int, ) -> None: + model = _unwrap_moe(model) assert is_mixture_of_experts(model) self.state = EplbState.from_mapping(