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 6, 00:45 UTC
[Feature] EPLB Support for GPU Model Runner v2 (#37488)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Wentao Ye committed
Mar 25, 2026 at 11:16 UTC
d7e93e13fb5c92243682466c1b3100b963445a17
3 files changed
+391
-6
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
new
+187
@@ -0,0 +1,187 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+
5
+from types import SimpleNamespace
6
+from typing import Any
7
+
8
+import torch
9
+
10
+from vllm.v1.worker.gpu import eplb_utils as eplb
11
+from vllm.v1.worker.gpu import model_runner as mrv2
12
+
13
+
14
+class FakeMemoryProfiler:
15
+ def __enter__(self):
16
+ self.consumed_memory = 0
17
+ return self
18
+
19
+ def __exit__(self, exc_type, exc, tb):
20
+ return False
21
+
22
+
23
+class FakeEplbState:
24
+ instances: list["FakeEplbState"] = []
25
+ from_mapping_kwargs: dict[str, Any] | None = None
26
+
27
+ def __init__(self, parallel_config: Any, device: torch.device):
28
+ self.parallel_config = parallel_config
29
+ self.device = device
30
+ self.add_model_calls: list[tuple[Any, Any]] = []
31
+ self.step_calls: list[tuple[bool, bool, bool]] = []
32
+ self.async_started = False
33
+ self.is_async = True
34
+ self.built_from_mapping = False
35
+ FakeEplbState.instances.append(self)
36
+
37
+ def add_model(self, model: Any, model_config: Any) -> None:
38
+ self.add_model_calls.append((model, model_config))
39
+
40
+ def step(self, is_dummy: bool, is_profile: bool, *, log_stats: bool) -> None:
41
+ self.step_calls.append((is_dummy, is_profile, log_stats))
42
+
43
+ def start_async_loop(self) -> None:
44
+ self.async_started = True
45
+
46
+ @classmethod
47
+ def from_mapping(cls, **kwargs: Any) -> "FakeEplbState":
48
+ cls.from_mapping_kwargs = kwargs
49
+ state = cls(kwargs["parallel_config"], kwargs["device"])
50
+ state.built_from_mapping = True
51
+ return state
52
+
53
+
54
+def _make_runner(**overrides: Any) -> Any:
55
+ runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)
56
+ runner.device = torch.device("cpu")
57
+ runner.model_config = SimpleNamespace(model="test-model")
58
+ runner.load_config = SimpleNamespace(load_format="hf")
59
+ runner.parallel_config = SimpleNamespace(
60
+ enable_eplb=True,
61
+ enable_elastic_ep=False,
62
+ eplb_config=SimpleNamespace(log_balancedness=True),
63
+ )
64
+ runner.vllm_config = SimpleNamespace(
65
+ load_config=runner.load_config,
66
+ model_config=runner.model_config,
67
+ )
68
+ runner.lora_config = None
69
+ runner.use_aux_hidden_state_outputs = False
70
+ runner.speculative_config = None
71
+ runner.speculator = None
72
+ runner.encoder_cache = None
73
+ runner.is_pooling_model = False
74
+ runner.is_last_pp_rank = True
75
+ runner.is_first_pp_rank = True
76
+ runner.max_num_reqs = 8
77
+ runner.max_num_tokens = 16
78
+ runner.decode_query_len = 1
79
+ runner.kv_connector = SimpleNamespace(set_disabled=lambda *_: None)
80
+ runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device)
81
+ runner.pooling_runner = None
82
+ runner.execute_model_state = None
83
+ for key, value in overrides.items():
84
+ setattr(runner, key, value)
85
+ return runner
86
+
87
+
88
+def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
89
+ FakeEplbState.instances.clear()
90
+ model = SimpleNamespace(is_moe=True)
91
+ prepared: list[object] = []
92
+
93
+ monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler)
94
+ monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
95
+ monkeypatch.setattr(
96
+ mrv2,
97
+ "get_model_loader",
98
+ lambda load_config: SimpleNamespace(load_model=lambda **_: model),
99
+ )
100
+ monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append)
101
+ monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state")
102
+ monkeypatch.setattr(
103
+ eplb,
104
+ "is_mixture_of_experts",
105
+ lambda loaded_model: getattr(loaded_model, "is_moe", False),
106
+ )
107
+
108
+ runner = _make_runner()
109
+ mrv2.GPUModelRunner.load_model(runner)
110
+
111
+ assert runner.model is model
112
+ assert runner.model_state == "model-state"
113
+ assert prepared == [model]
114
+ assert runner.eplb_state is not None
115
+ assert runner.eplb_state.add_model_calls == [(model, runner.model_config)]
116
+ assert runner.eplb_state.async_started is True
117
+
118
+
119
+def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
120
+ FakeEplbState.instances.clear()
121
+ model = SimpleNamespace(is_moe=True)
122
+ prepared: list[object] = []
123
+
124
+ monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler)
125
+ monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
126
+ monkeypatch.setattr(
127
+ mrv2,
128
+ "get_model_loader",
129
+ lambda load_config: SimpleNamespace(load_model=lambda **_: model),
130
+ )
131
+ monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append)
132
+ monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state")
133
+ monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
134
+
135
+ runner = _make_runner()
136
+ mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True)
137
+
138
+ assert runner.load_config.load_format == "dummy"
139
+ assert prepared == []
140
+ assert runner.eplb_state is not None
141
+ assert runner.eplb_state.add_model_calls == []
142
+ assert runner.eplb_state.async_started is False
143
+
144
+
145
+def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch):
146
+ FakeEplbState.instances.clear()
147
+ FakeEplbState.from_mapping_kwargs = None
148
+ monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
149
+ monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
150
+
151
+ runner = _make_runner(model=SimpleNamespace(is_moe=True))
152
+ mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)
153
+ mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping, 2)
154
+
155
+ assert runner.eplb_state is not None
156
+ assert runner.eplb_state.built_from_mapping is True
157
+ assert FakeEplbState.from_mapping_kwargs is not None
158
+ assert FakeEplbState.from_mapping_kwargs["expanded_physical_to_logical"] is mapping
159
+ assert FakeEplbState.from_mapping_kwargs["num_valid_physical_experts"] == 2
160
+
161
+
162
+def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
163
+ events = []
164
+ runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0)
165
+ runner.execute_model_state = SimpleNamespace(
166
+ input_batch=SimpleNamespace(num_reqs=2),
167
+ attn_metadata=None,
168
+ slot_mappings_by_layer=None,
169
+ hidden_states=None,
170
+ aux_hidden_states=None,
171
+ kv_connector_output=None,
172
+ num_tokens_across_dp=None,
173
+ )
174
+ runner.postprocess = lambda *args, **kwargs: events.append("postprocess")
175
+ runner.eplb.step = lambda *args, **kwargs: events.append("eplb")
176
+ monkeypatch.setattr(
177
+ mrv2,
178
+ "pp_receive",
179
+ lambda *args, **kwargs: (
180
+ torch.zeros((2, 1), dtype=torch.long),
181
+ torch.ones(2, dtype=torch.int32),
182
+ torch.zeros(2, dtype=torch.int32),
183
+ ),
184
+ )
185
+
186
+ assert mrv2.GPUModelRunner.sample_tokens(runner, None) is None
187
+ assert events == ["postprocess", "eplb"]
vllm/v1/worker/gpu/eplb_utils.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable
5
+from functools import wraps
6
+from typing import Any
7
+
8
+import torch
9
+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
14
+
15
+logger = init_logger(__name__)
16
+
17
+
18
+def step_eplb_after(*, is_dummy: bool = False) -> Callable:
19
+ """Step EPLB after a model runner method completes successfully."""
20
+
21
+ def decorator(fn: Callable) -> Callable:
22
+ @wraps(fn)
23
+ def wrapper(self: Any, *args, **kwargs) -> Any:
24
+ result = fn(self, *args, **kwargs)
25
+ if kwargs.get("skip_eplb", False):
26
+ return result
27
+
28
+ is_profile = kwargs.get("is_profile", False) if is_dummy else False
29
+ self.eplb.step(is_dummy=is_dummy, is_profile=is_profile)
30
+ return result
31
+
32
+ return wrapper
33
+
34
+ return decorator
35
+
36
+
37
+class EPLBController:
38
+ def __init__(self, parallel_config: Any, device: torch.device):
39
+ self.parallel_config = parallel_config
40
+ self.device = device
41
+ self.state: EplbState | None = None
42
+ self.suppressed = False
43
+ self._has_registered_models = False
44
+
45
+ def prepare_load(self) -> None:
46
+ self.state = None
47
+ self._has_registered_models = False
48
+ if self.parallel_config.enable_eplb:
49
+ self.state = EplbState(self.parallel_config, self.device)
50
+
51
+ def maybe_register_speculator(
52
+ self,
53
+ speculator: Any | None,
54
+ speculative_config: Any | None,
55
+ load_dummy_weights: bool,
56
+ ) -> bool:
57
+ # if speculator is a moe model, add it to eplb
58
+ if (
59
+ speculator is None
60
+ or not hasattr(speculator, "model")
61
+ or not self.parallel_config.enable_eplb
62
+ or load_dummy_weights
63
+ ):
64
+ return False
65
+
66
+ draft_model = speculator.model
67
+ if not is_mixture_of_experts(draft_model):
68
+ return False
69
+
70
+ assert not self.parallel_config.enable_elastic_ep, (
71
+ "Elastic EP is not supported with draft model."
72
+ )
73
+ assert speculative_config is not None
74
+ assert speculative_config.draft_model_config is not None
75
+ assert self.state is not None
76
+ self.state.add_model(
77
+ draft_model,
78
+ speculative_config.draft_model_config,
79
+ )
80
+ self._has_registered_models = True
81
+ return True
82
+
83
+ def maybe_register_model(
84
+ self,
85
+ model: nn.Module,
86
+ model_config: Any,
87
+ load_dummy_weights: bool,
88
+ ) -> bool:
89
+ if not self.parallel_config.enable_eplb or load_dummy_weights:
90
+ return False
91
+
92
+ if not is_mixture_of_experts(model):
93
+ return False
94
+
95
+ logger.info_once(
96
+ "EPLB is enabled for model %s.", model_config.model, scope="local"
97
+ )
98
+ assert self.state is not None
99
+ self.state.add_model(model, model_config)
100
+ self._has_registered_models = True
101
+ return True
102
+
103
+ def maybe_start_async_loop(self, eplb_models_added: bool) -> None:
104
+ if eplb_models_added and self.state is not None and self.state.is_async:
105
+ self.state.start_async_loop()
106
+
107
+ def step(
108
+ self,
109
+ is_dummy: bool = False,
110
+ is_profile: bool = False,
111
+ ) -> None:
112
+ if (
113
+ not self.parallel_config.enable_eplb
114
+ or self.suppressed
115
+ or self.state is None
116
+ or not self._has_registered_models
117
+ ):
118
+ return
119
+
120
+ self.state.step(
121
+ is_dummy,
122
+ is_profile,
123
+ log_stats=self.parallel_config.eplb_config.log_balancedness,
124
+ )
125
+
126
+ def setup_from_mapping(
127
+ self,
128
+ model: nn.Module,
129
+ model_config: Any,
130
+ expanded_physical_to_logical: torch.Tensor,
131
+ old_num_physical_experts: int,
132
+ ) -> None:
133
+ assert is_mixture_of_experts(model)
134
+
135
+ self.state = EplbState.from_mapping(
136
+ model=model,
137
+ model_config=model_config,
138
+ device=self.device,
139
+ parallel_config=self.parallel_config,
140
+ expanded_physical_to_logical=expanded_physical_to_logical,
141
+ num_valid_physical_experts=old_num_physical_experts,
142
+ )
143
+ self._has_registered_models = True
vllm/v1/worker/gpu/model_runner.py
+61
-6
@@ -62,6 +62,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
62
get_uniform_token_count,
63
)
64
from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
65
+from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
66
from vllm.v1.worker.gpu.input_batch import (
67
InputBatch,
68
InputBuffers,
@@ -244,6 +245,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
245
# For transferring state from execute_model to subsequent sample_tokens call.
246
self.execute_model_state: ExecuteModelState | None = None
247
248
+ # Expert parallelism load balancer.
249
+ self.eplb = EPLBController(self.parallel_config, self.device)
250
+
251
def update_max_model_len(self, max_model_len: int) -> None:
252
self.max_model_len = max_model_len
253
self.req_states.max_model_len = max_model_len
@@ -259,8 +263,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
263
tasks.extend(PoolingRunner.get_supported_tasks(self.model))
264
return tuple(tasks)
265
262
- def load_model(self, *args, **kwargs) -> None:
266
+ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
267
time_before_load = time.perf_counter()
268
+ if load_dummy_weights:
269
+ self.load_config.load_format = "dummy"
270
+ self.eplb.prepare_load()
271
+ eplb_models_added = False
272
with DeviceMemoryProfiler() as m:
273
model_loader = get_model_loader(self.vllm_config.load_config)
274
logger.info("Loading model from scratch...")
@@ -278,6 +286,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
286
set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
287
if self.speculator is not None:
288
self.speculator.load_model(self.model)
289
+ eplb_models_added = self.eplb.maybe_register_speculator(
290
+ self.speculator, self.speculative_config, load_dummy_weights
291
+ )
292
time_after_load = time.perf_counter()
293
294
self.model_memory_usage = m.consumed_memory
@@ -287,9 +298,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
298
time_after_load - time_before_load,
299
)
300
290
- prepare_communication_buffer_for_model(self.model)
291
- if self.speculator is not None:
292
- prepare_communication_buffer_for_model(self.speculator.model)
301
+ if not load_dummy_weights:
302
+ prepare_communication_buffer_for_model(self.model)
303
+ if self.speculator is not None:
304
+ prepare_communication_buffer_for_model(self.speculator.model)
305
306
# Initialize the components that require the model.
307
self.model_state = init_model_state(
@@ -297,6 +309,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
309
)
310
if self.is_pooling_model and self.is_last_pp_rank:
311
self.pooling_runner = PoolingRunner(self.model)
312
+ eplb_models_added |= self.eplb.maybe_register_model(
313
+ self.model,
314
+ self.model_config,
315
+ load_dummy_weights,
316
+ )
317
+ self.eplb.maybe_start_async_loop(eplb_models_added)
318
319
if not self.is_first_pp_rank:
320
# For non-first PP ranks, create intermediate tensors sized
@@ -372,12 +390,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
390
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
391
392
@torch.inference_mode()
393
+ @step_eplb_after(is_dummy=True)
394
def _dummy_run(
395
self,
396
num_tokens: int,
397
*args,
398
skip_attn: bool = True,
399
uniform_decode: bool = False,
400
+ skip_eplb: bool = False,
401
+ is_profile: bool = False,
402
**kwargs,
403
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
404
# Create a dummy scheduler output.
@@ -493,7 +514,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
514
@torch.inference_mode()
515
def profile_run(self) -> None:
516
hidden_states, sample_hidden_states = self._dummy_run(
496
- self.max_num_tokens, skip_attn=True
517
+ self.max_num_tokens, skip_attn=True, is_profile=True
518
)
519
520
# Only run sampler/pooler on last PP rank (non-last ranks return None).
@@ -1090,6 +1111,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1111
return None
1112
1113
@torch.inference_mode()
1114
+ @step_eplb_after()
1115
def sample_tokens(
1116
self, grammar_output: GrammarOutput | None
1117
) -> AsyncOutput | ModelRunnerOutput | None:
@@ -1211,6 +1233,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1233
return self.draft_tokens_handler.get_draft_tokens()
1234
1235
@torch.inference_mode()
1236
+ @step_eplb_after()
1237
def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
1238
if self.execute_model_state is None:
1239
# The prior execute_model call must have failed.
@@ -1229,7 +1252,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1252
pooler_output, is_valid = self.pooling_runner.pool(
1253
hidden_states, input_batch, self.req_states
1254
)
1232
- self.postprocess_pool(input_batch)
1255
1256
# Build the model runner output.
1257
model_runner_output = ModelRunnerOutput(
@@ -1245,6 +1267,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1267
copy_stream=self.output_copy_stream,
1268
copy_event=self.output_copy_event,
1269
)
1270
+
1271
+ self.postprocess_pool(input_batch)
1272
if self.use_async_scheduling:
1273
return async_output
1274
return async_output.get_output()
@@ -1265,6 +1289,37 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1289
computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
1290
)
1291
1292
+ ########### EPLB methods start ###########
1293
+ @property
1294
+ def eplb_state(self):
1295
+ return self.eplb.state
1296
+
1297
+ @eplb_state.setter
1298
+ def eplb_state(self, state) -> None:
1299
+ self.eplb.state = state
1300
+
1301
+ @property
1302
+ def eep_eplb_suppressed(self) -> bool:
1303
+ return self.eplb.suppressed
1304
+
1305
+ @eep_eplb_suppressed.setter
1306
+ def eep_eplb_suppressed(self, suppressed: bool) -> None:
1307
+ self.eplb.suppressed = suppressed
1308
+
1309
+ def setup_eplb_from_mapping(
1310
+ self,
1311
+ expanded_physical_to_logical: torch.Tensor,
1312
+ old_num_physical_experts: int,
1313
+ ) -> None:
1314
+ self.eplb.setup_from_mapping(
1315
+ self.model,
1316
+ self.model_config,
1317
+ expanded_physical_to_logical,
1318
+ old_num_physical_experts,
1319
+ )
1320
+
1321
+ ########### EPLB methods end ###########
1322
+
1323
1324
class ExecuteModelState(NamedTuple):
1325
input_batch: InputBatch