[Bugfix] Unwrap VLM wrappers for EPLB on Model Runner V2 (#42706)

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(