A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 13:00 UTC
[Bugfix] Unwrap VLM wrappers for EPLB on Model Runner V2 (#42706)
Keyi Li committed
May 15, 2026 at 16:20 UTC
d0921bafeff9bbe7a7b4efef6371700e69224702
1 file changed
+16
-1
vllm/v1/worker/gpu/eplb_utils.py
+16
-1
@@ -10,11 +10,24 @@ import torch.nn as nn
10
11
from vllm.distributed.eplb.eplb_state import EplbState
12
from vllm.logger import init_logger
13
-from vllm.model_executor.models.interfaces import is_mixture_of_experts
13
+from vllm.model_executor.models.interfaces import (
14
+ SupportsMultiModal,
15
+ is_mixture_of_experts,
16
+)
17
18
logger = init_logger(__name__)
19
20
21
+def _unwrap_moe(model: nn.Module) -> nn.Module:
22
+ # VLM wrappers (e.g. KimiK25ForConditionalGeneration) hold the MoE
23
+ # language model under `.language_model` but don't implement
24
+ # MixtureOfExperts themselves. Mirror the V1 path
25
+ # (see vllm/v1/worker/gpu_model_runner.py, PR #39805).
26
+ if not is_mixture_of_experts(model) and isinstance(model, SupportsMultiModal):
27
+ return model.get_language_model()
28
+ return model
29
+
30
+
31
def step_eplb_after(*, is_dummy: bool = False) -> Callable:
32
"""Step EPLB after a model runner method completes successfully."""
33
@@ -89,6 +102,7 @@ class EPLBController:
102
if not self.parallel_config.enable_eplb or load_dummy_weights:
103
return False
104
105
+ model = _unwrap_moe(model)
106
if not is_mixture_of_experts(model):
107
return False
108
@@ -128,6 +142,7 @@ class EPLBController:
142
expanded_physical_to_logical: torch.Tensor,
143
old_num_physical_experts: int,
144
) -> None:
145
+ model = _unwrap_moe(model)
146
assert is_mixture_of_experts(model)
147
148
self.state = EplbState.from_mapping(