From d0921bafeff9bbe7a7b4efef6371700e69224702 Mon Sep 17 00:00:00 2001
From: Keyi Li <94494390+JasonKeyiL@users.noreply.github.com>
Date: Fri, 15 May 2026 16:20:33 -0700
Subject: [PATCH] [Bugfix] Unwrap VLM wrappers for EPLB on Model Runner V2
 (#42706)

---
 vllm/v1/worker/gpu/eplb_utils.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

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(