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, 17:45 UTC
[Bugfix] Fix EPLB initialization for VLM wrapper models (#39805)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Roy Wang committed
May 14, 2026 at 10:26 UTC
77e1421a68d625bb45ac1d86023be8e555aefeec
1 file changed
+20
-8
vllm/v1/worker/gpu_model_runner.py
+20
-8
@@ -73,6 +73,7 @@ from vllm.model_executor.model_loader.reload import (
73
initialize_layerwise_reload,
74
)
75
from vllm.model_executor.models.interfaces import (
76
+ MixtureOfExperts,
77
MultiModalEmbeddings,
78
SupportsMRoPE,
79
SupportsMultiModal,
@@ -489,6 +490,7 @@ class GPUModelRunner(
490
self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
491
492
self.eplb_state: EplbState | None = None
493
+ self._moe_model: MixtureOfExperts | None = None
494
# NOTE(yongji): flag to temporarily disable EPLB during scaling up/down
495
self.eep_eplb_suppressed = False
496
"""
@@ -3186,8 +3188,7 @@ class GPUModelRunner(
3188
return
3189
3190
assert self.eplb_state is not None
3189
- model = self.get_model()
3190
- assert is_mixture_of_experts(model)
3191
+ assert self._moe_model is not None
3192
self.eplb_state.step(
3193
is_dummy,
3194
is_profile,
@@ -3199,11 +3200,10 @@ class GPUModelRunner(
3200
expanded_physical_to_logical: torch.Tensor,
3201
old_num_physical_experts: int,
3202
) -> None:
3202
- model = self.get_model()
3203
- assert is_mixture_of_experts(model)
3203
+ assert self._moe_model is not None
3204
3205
self.eplb_state = EplbState.from_mapping(
3206
- model=model,
3206
+ model=self._moe_model,
3207
model_config=self.model_config,
3208
device=self.device,
3209
parallel_config=self.parallel_config,
@@ -4996,8 +4996,20 @@ class GPUModelRunner(
4996
4997
self.model.set_aux_hidden_state_layers(aux_layers)
4998
4999
+ # Resolve the MoE model, unwrapping VLM wrappers if needed.
5000
+ # VLM models (e.g. KimiK25ForConditionalGeneration) wrap the
5001
+ # actual MoE language model but don't implement
5002
+ # MixtureOfExperts themselves.
5003
+ moe_candidate = self.model
5004
+ if not is_mixture_of_experts(moe_candidate) and isinstance(
5005
+ moe_candidate, SupportsMultiModal
5006
+ ):
5007
+ moe_candidate = moe_candidate.get_language_model()
5008
+ if is_mixture_of_experts(moe_candidate):
5009
+ self._moe_model = moe_candidate
5010
+
5011
if (
5000
- is_mixture_of_experts(self.model)
5012
+ self._moe_model is not None
5013
and self.parallel_config.enable_eplb
5014
and not load_dummy_weights
5015
):
@@ -5007,7 +5019,7 @@ class GPUModelRunner(
5019
)
5020
assert self.eplb_state is not None
5021
self.eplb_state.add_model(
5010
- self.model,
5022
+ self._moe_model,
5023
self.model_config,
5024
)
5025
eplb_models += 1
@@ -5047,7 +5059,7 @@ class GPUModelRunner(
5059
) # Temporary hack for dynamic res video w/o support for bs>1 yet
5060
5061
if (
5050
- is_mixture_of_experts(self.model)
5062
+ self._moe_model is not None
5063
and self.parallel_config.enable_eplb
5064
and not load_dummy_weights
5065
and self.eplb_state is not None