@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 171d1b9ef6

[Elastic EP] Support Model Runner V2 (#53934)

Signed-off-by: almogtavor <almogtavor@gmail.com> Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com>

Almog Tavor committed Sep 29, 2026 at 02:20 UTC 171d1b9ef6fdfb86a45d2c540b52535e4d5a660f

12 files changed +171 -80

tests/v1/cudagraph/test_cudagraph_manager.py

+3 -1

@@ -224,7 +224,9 @@ def test_speculator_capture_preserves_decode_query_bounds(

224 )

225 manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)

226 manager.max_num_reqs = num_reqs

227 - manager.dp_size = 1

227 + manager.vllm_config = SimpleNamespace(

228 + parallel_config=SimpleNamespace(data_parallel_size=1)

229 + )

230 buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))

231 block_tables = MagicMock()

232 block_tables.cp_size = 1

tests/v1/engine/test_engine_core_client.py

@@ -215,6 +215,10 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM

215 client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]

216 client.eng_start_index = 0

217 client._kv_event_sources = {}

218 + client.vllm_config = SimpleNamespace(

219 + use_v2_model_runner=True,

220 + parallel_config=SimpleNamespace(all2all_backend="allgather_reducescatter"),

221 + )

222 return client

223

224

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

+12 -21

@@ -24,7 +24,6 @@ class FakeMemoryProfiler:

25 class FakeEplbState:

26 instances: list["FakeEplbState"] = []

27 - from_mapping_kwargs: dict[str, Any] | None = None

28 def __init__(self, parallel_config: Any, device: torch.device):

29 self.parallel_config = parallel_config

@@ -33,7 +32,7 @@ class FakeEplbState:

32 self.step_calls: list[tuple[bool, bool, bool]] = []

33 self.async_started = False

34 self.is_async = True

36 - self.built_from_mapping = False

35 + self.update_mapping_calls: list[tuple[Any, torch.Tensor]] = []

36 FakeEplbState.instances.append(self)

38 def add_model(self, model: Any, model_config: Any) -> None:

@@ -45,12 +44,8 @@ class FakeEplbState:

44 def start_async_loop(self) -> None:

45 self.async_started = True

48 - @classmethod

49 - def from_mapping(cls, **kwargs: Any) -> "FakeEplbState":

50 - cls.from_mapping_kwargs = kwargs

51 - state = cls(kwargs["parallel_config"], kwargs["device"])

52 - state.built_from_mapping = True

53 - return state

47 + def update_mapping(self, model_config: Any, mapping: torch.Tensor) -> None:

48 + self.update_mapping_calls.append((model_config, mapping))

51 def _make_runner(**overrides: Any) -> Any:

@@ -136,7 +131,7 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):

131 assert runner.eplb_state.async_started is True

132

133

139 -def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):

134 +def test_v2_load_model_with_dummy_weights_defers_eplb_async_loop(monkeypatch):

135 FakeEplbState.instances.clear()

136 model = SimpleNamespace(is_moe=True)

137

@@ -159,24 +154,20 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):

154

155 assert runner.load_config.load_format == "dummy"

156 assert runner.eplb_state is not None

162 - assert runner.eplb_state.add_model_calls == []

157 + # A scaling-up worker registers so it can join the EPLB communicator.

158 + assert runner.eplb_state.add_model_calls == [(model, runner.model_config)]

159 assert runner.eplb_state.async_started is False

160

161

166 -def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch):

167 - FakeEplbState.instances.clear()

168 - FakeEplbState.from_mapping_kwargs = None

169 - monkeypatch.setattr(eplb, "EplbState", FakeEplbState)

170 - monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model)

171 -

172 - runner = _make_runner(model=SimpleNamespace(is_moe=True))

162 +def test_v2_setup_eplb_from_mapping_updates_state_in_place():

163 + runner = _make_runner()

164 + state = FakeEplbState(runner.parallel_config, runner.device)

165 + runner.eplb_state = state

166 mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)

167 mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping)

168

176 - assert runner.eplb_state is not None

177 - assert runner.eplb_state.built_from_mapping is True

178 - assert FakeEplbState.from_mapping_kwargs is not None

179 - assert FakeEplbState.from_mapping_kwargs["expanded_physical_to_logical"] is mapping

169 + assert runner.eplb_state is state

170 + assert state.update_mapping_calls == [(runner.model_config, mapping)]

171

172

173 def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):

vllm/config/vllm.py

-3

@@ -3086,9 +3086,6 @@ class VllmConfig:

3086 if self.parallel_config.use_ubatching:

3087 unsupported.extend(self._get_dbo_unsupported_features())

3088

3089 - if self.parallel_config.enable_elastic_ep:

3090 - unsupported.append("elastic expert parallelism")

3091 -

3089 if self.cache_config.mamba_cache_mode == "all":

3090 unsupported.append("mamba cache mode 'all'")

3091

vllm/distributed/elastic_ep/elastic_execute.py

+33 -22

@@ -54,6 +54,8 @@ from vllm.platforms import current_platform

54 from vllm.utils import is_moe_layer

55 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

56 from vllm.v1.worker.dp_utils import skip_dp_coordination

57 +from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager

58 +from vllm.v1.worker.gpu.eplb_utils import preserve_serving_state

59 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper

60 from vllm.v1.worker.workspace import lock_workspace, unlock_workspace

@@ -171,6 +173,17 @@ class ElasticEPScalingExecutor:

173 self.worker.vllm_config.parallel_config,

174 )

175

176 + @contextmanager

177 + def _suppress_eplb(self) -> Iterator[None]:

178 + # A rank warming up alone must not advance the EPLB step counter.

179 + runner = self.worker.model_runner

180 + was_suppressed = runner.eep_eplb_suppressed

181 + runner.eep_eplb_suppressed = True

182 + try:

183 + yield

184 + finally:

185 + runner.eep_eplb_suppressed = was_suppressed

186 +

187 @contextmanager

188 def _disable_flashinfer_autotune(self) -> Iterator[None]:

189 kernel_config = self.worker.vllm_config.kernel_config

@@ -421,7 +434,17 @@ class ElasticEPScalingExecutor:

434 self._staged_moe_quant_methods.clear()

435

436 def _release_cuda_graphs(self) -> None:

424 - if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):

437 + manager = getattr(self.worker.model_runner, "cudagraph_manager", None)

438 + speculator = getattr(self.worker.model_runner, "speculator", None)

439 + if manager is not None:

440 + # MRV2 captures through CudaGraphManager instead of wrapping the

441 + # model, so neither wrapper branch below ever fires.

442 + manager.release_graphs()

443 + for attr in vars(speculator).values() if speculator else ():

444 + if isinstance(attr, CudaGraphManager):

445 + attr.release_graphs()

446 +

447 + elif isinstance(self.worker.model_runner.model, CUDAGraphWrapper):

448 wrapper = self.worker.model_runner.model

449 wrapper.concrete_cudagraph_entries = {}

450

@@ -431,6 +454,8 @@ class ElasticEPScalingExecutor:

454 torch.compiler.reset()

455 with set_current_vllm_config(self.worker.vllm_config):

456 reset_compile_wrapper(self.worker.model_runner.get_model())

457 + if speculator is not None:

458 + reset_compile_wrapper(speculator.model)

459

460 gc.collect()

461 torch.accelerator.synchronize()

@@ -675,42 +700,28 @@ class ElasticEPScalingExecutor:

700 self.warm_and_capture()

701

702 def warm_and_capture(self) -> None:

678 - # Save and clear block tables so the dummy MoE forward doesn't

679 - # write dummy slot mappings into real KV-cache blocks.

680 - multi_block_table = self.worker.model_runner.input_batch.block_table

681 - saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []

682 - for bt in multi_block_table.block_tables:

683 - saved_block_tables.append(

684 - (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())

685 - )

686 - multi_block_table.clear()

687 -

703 # _ensure_workspace_size allocates a fresh tensor on grow, leaving

704 # any captured CUDA graph with a stale data pointer; drop graphs

705 # before re-warm so captures realign with the resized buffer.

706 self._release_cuda_graphs()

707 unlock_workspace()

708

694 - # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model

695 - # alone only exercises cudagraph-capture sizes and can leave the

696 - # workspace too small for post-reshuffle routing. Use _dummy_run

697 - # directly with skip_eplb=True so dummy routing doesn't pollute the

698 - # just-rebalanced EPLB stats.

709 runner = self.worker.model_runner

710 all2all_manager = get_ep_all2all_manager()

711 reuse_kernel = self._can_reuse_fused_moe_kernel()

712 + serving_state = (

713 + preserve_serving_state(runner)

714 + if self.worker.use_v2_model_runner

715 + else runner.preserve_serving_state()

716 + )

717 with (

718 skip_dp_coordination() if reuse_kernel else nullcontext(),

719 all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(),

720 self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(),

721 + self._suppress_eplb(),

722 + serving_state,

723 ):

724 runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)

725 self.worker.compile_or_warm_up_model()

726

727 lock_workspace()

711 -

712 - for bt, (saved_gpu, saved_cpu) in zip(

713 - multi_block_table.block_tables, saved_block_tables

714 - ):

715 - bt.block_table.gpu.copy_(saved_gpu)

716 - bt.block_table.cpu.copy_(saved_cpu)

vllm/v1/engine/core_client.py

+19 -2

@@ -1893,13 +1893,29 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1893 await asyncio.to_thread(self._coord_store.wait, ready_keys)

1894 logger.info("[Elastic EP] Successfully started new engines")

1895

1896 + def _eep_commit_pause_mode(self) -> PauseMode:

1897 + from vllm.distributed.elastic_ep.elastic_execute import (

1898 + can_reuse_fused_moe_kernel,

1899 + )

1900 +

1901 + # MRV2 re-warms through the request pool, so running requests must

1902 + # finish first while new ones stay queued in the scheduler.

1903 + parallel_config = self.vllm_config.parallel_config

1904 + if self.vllm_config.use_v2_model_runner and not can_reuse_fused_moe_kernel(

1905 + parallel_config

1906 + ):

1907 + return "wait"

1908 + return "keep"

1909 +

1910 async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None:

1911 new_core_engines = [

1912 rank.to_bytes(2, "little")

1913 for rank in range(len(self.core_engines), new_data_parallel_size)

1914 ]

1915

1902 - await self.pause_scheduler_async(mode="keep", clear_cache=False)

1916 + await self.pause_scheduler_async(

1917 + mode=self._eep_commit_pause_mode(), clear_cache=False

1918 + )

1919 wait_future = self._eep_wait_for_setup_switch_complete()

1920 finish_futures = [

1921 asyncio.create_task(

@@ -1971,7 +1987,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1987 for rank in range(new_data_parallel_size, cur_data_parallel_size):

1988 self._kv_event_sources.pop(rank, None)

1989 removed_dp_size = cur_data_parallel_size - new_data_parallel_size

1974 - pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size

1990 + eep_mode = self._eep_commit_pause_mode()

1991 + pause_modes = [eep_mode] * new_data_parallel_size + ["abort"] * removed_dp_size

1992 pause_futures = [

1993 self._call_utility_async("pause_scheduler", mode, False, engine=engine)

1994 for mode, engine in zip(pause_modes, old_core_engines)

vllm/v1/worker/gpu/block_table.py

@@ -15,6 +15,9 @@ from vllm.v1.worker.gpu.buffer_utils import (

17 class BlockTables:

18 + # Set by the elastic EP warmup so KV writes land in the null block.

19 + redirect_writes_to_null_block = False

20 +

21 def __init__(

22 self,

23 block_sizes: list[int],

@@ -126,6 +129,8 @@ class BlockTables:

129 bpk = self.blocks_per_kv_block[i]

130 if bpk > 1:

131 block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]

132 + if self.redirect_writes_to_null_block:

133 + block_ids = [0] * len(block_ids)

134 end = start + len(block_ids)

135 row_capacity = self.block_tables[i].gpu.shape[1]

136 if end > row_capacity:

vllm/v1/worker/gpu/cudagraph_utils.py

+17 -1

@@ -161,7 +161,6 @@ class CudaGraphManager:

161 # DBO supports FULL CUDA graphs only.

162 self.ubatch_runner = ubatch_runner

163

164 - self.dp_size = vllm_config.parallel_config.data_parallel_size

164 self.tp_size = vllm_config.parallel_config.tensor_parallel_size

165 self.is_first_pp_rank = get_pp_group().is_first_rank

166 self.is_last_pp_rank = get_pp_group().is_last_rank

@@ -385,6 +384,23 @@ class CudaGraphManager:

384 self._candidates.setdefault(key, []).extend(matching)

385 current_range_start = num_tokens + 1

386

387 + @property

388 + def dp_size(self) -> int:

389 + # Not cached: elastic EP rewrites parallel_config in place on scale.

390 + return self.vllm_config.parallel_config.data_parallel_size

391 +

392 + def release_graphs(self) -> None:

393 + """Drop the captured graphs so a later capture() can refill them.

394 +

395 + Elastic EP reallocates the MoE workspace when it grows, which leaves

396 + every captured graph holding a stale data pointer. `_capture_descs` is

397 + kept, so `needs_capture()` still reports the work to redo.

398 + """

399 + self.graphs.clear()

400 + self._graphs_captured = False

401 + if self.breakable_cg_runner is not None:

402 + BreakableCUDAGraphWrapper.clear_all_graphs()

403 +

404 def needs_capture(self) -> bool:

405 return len(self._capture_descs) > 0

406

vllm/v1/worker/gpu/eplb_utils.py

+41 -20

@@ -1,9 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from collections.abc import Callable

4 +from collections.abc import Callable, Iterator

5 +from contextlib import contextmanager

6 from functools import wraps

6 -from typing import Any

7 +from typing import TYPE_CHECKING, Any

9 import torch

10 import torch.nn as nn

@@ -15,6 +16,9 @@ from vllm.model_executor.models.interfaces import (

16 get_mixture_of_experts_model,

17 )

19 +if TYPE_CHECKING:

20 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner

21 +

22 logger = init_logger(__name__)

@@ -37,6 +41,27 @@ def step_eplb_after(*, is_dummy: bool = False) -> Callable:

41 return decorator

44 +@contextmanager

45 +def preserve_serving_state(model_runner: "GPUModelRunner") -> Iterator[None]:

46 + """Keep the elastic EP warmup out of the request pool and the KV cache."""

47 + # After the drain only parked streaming sessions still hold a slot, and

48 + # those are re-added from NewRequestData on their next chunk.

49 + _remove_all_requests(model_runner)

50 + model_runner.block_tables.redirect_writes_to_null_block = True

51 + try:

52 + yield

53 + finally:

54 + _remove_all_requests(model_runner)

55 + model_runner.block_tables.redirect_writes_to_null_block = False

56 + if model_runner.kv_block_zeroer is not None:

57 + model_runner.kv_block_zeroer.zero_block_ids([0])

58 +

59 +

60 +def _remove_all_requests(model_runner: "GPUModelRunner") -> None:

61 + for req_id in list(model_runner.req_states.req_id_to_index):

62 + model_runner._remove_request(req_id)

63 +

64 +

65 class EPLBController:

66 def __init__(self, parallel_config: Any, device: torch.device):

67 self.parallel_config = parallel_config

@@ -82,10 +107,8 @@ class EPLBController:

107 self._has_registered_models = True

108 return True

109

85 - def maybe_register_model(

86 - self, model: nn.Module, model_config: Any, load_dummy_weights: bool

87 - ) -> bool:

88 - if not self.parallel_config.enable_eplb or load_dummy_weights:

110 + def maybe_register_model(self, model: nn.Module, model_config: ModelConfig) -> bool:

111 + if not self.parallel_config.enable_eplb:

112 return False

113

114 moe_model = get_mixture_of_experts_model(model)

@@ -100,7 +123,13 @@ class EPLBController:

123 self._has_registered_models = True

124 return True

125

103 - def maybe_start_async_loop(self, eplb_models_added: bool) -> None:

126 + def maybe_start_async_loop(

127 + self, eplb_models_added: bool, load_dummy_weights: bool = False

128 + ) -> None:

129 + # A scaling-up worker registers to join the communicator handshake, but

130 + # its groups are not live yet, so _perform_eplb_reshuffle() starts this.

131 + if load_dummy_weights:

132 + return

133 if eplb_models_added and self.state is not None and self.state.is_async:

134 self.state.start_async_loop()

135

@@ -128,18 +157,10 @@ class EPLBController:

157

158 def setup_from_mapping(

159 self,

131 - model: nn.Module,

132 - model_config: Any,

160 + model_config: ModelConfig,

161 expanded_physical_to_logical: torch.Tensor,

162 ) -> None:

135 - moe_model = get_mixture_of_experts_model(model)

136 - assert moe_model is not None

137 -

138 - self.state = EplbState.from_mapping(

139 - model=moe_model,

140 - model_config=model_config,

141 - device=self.device,

142 - parallel_config=self.parallel_config,

143 - expanded_physical_to_logical=expanded_physical_to_logical,

144 - )

145 - self._has_registered_models = True

163 + # Update in place: a fresh EplbState would build a second communicator

164 + # over a group the other ranks have already left.

165 + assert self.state is not None

166 + self.state.update_mapping(model_config, expanded_physical_to_logical)

vllm/v1/worker/gpu/model_runner.py

+13 -9

@@ -234,10 +234,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

234 # Persistent buffer for intermediate tensors (non-first PP ranks).

235 self.intermediate_tensors: IntermediateTensors | None = None

236

237 - # Data parallelism.

238 - self.dp_size = self.parallel_config.data_parallel_size

239 - self.dp_rank = self.parallel_config.data_parallel_rank

240 -

237 # Dual batch overlap. Created in initialize_kv_cache(), once everything

238 # it runs the microbatched forward with exists.

239 self.ubatch_runner: UBatchRunner | None = None

@@ -509,9 +505,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

505 if self.is_pooling_model and self.is_last_pp_rank:

506 self.pooling_runner = PoolingRunner(self.model, self.vllm_config)

507 eplb_models_added |= self.eplb.maybe_register_model(

512 - self.model, self.model_config, load_dummy_weights

508 + self.model, self.model_config

509 )

514 - self.eplb.maybe_start_async_loop(eplb_models_added)

510 + self.eplb.maybe_start_async_loop(eplb_models_added, load_dummy_weights)

511

512 if not self.is_first_pp_rank:

513 # For non-first PP ranks, create intermediate tensors sized

@@ -526,6 +522,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

522

523 get_offloader().post_init()

524

525 + @property

526 + def dp_size(self) -> int:

527 + # Not cached: elastic EP rewrites parallel_config in place on reconfigure.

528 + return self.parallel_config.data_parallel_size

529 +

530 + @property

531 + def dp_rank(self) -> int:

532 + # Not cached: elastic EP rewrites parallel_config in place on reconfigure.

533 + return self.parallel_config.data_parallel_rank

534 +

535 def get_model(self) -> nn.Module:

536 return self.model

537

@@ -2337,9 +2343,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2343 def setup_eplb_from_mapping(

2344 self, expanded_physical_to_logical: torch.Tensor

2345 ) -> None:

2340 - self.eplb.setup_from_mapping(

2341 - self.model, self.model_config, expanded_physical_to_logical

2342 - )

2346 + self.eplb.setup_from_mapping(self.model_config, expanded_physical_to_logical)

2347

2348 ########### EPLB methods end ###########

2349

vllm/v1/worker/gpu/warmup.py

+7 -1

@@ -294,7 +294,11 @@ def _warmup_kernels(

294 model_runner.scheduler_config.max_num_batched_tokens

295 // max(prompt_len, decode_query_len),

296 )

297 - if max_blocks_per_req > 0:

297 + block_tables = getattr(model_runner, "block_tables", None)

298 + null_blocks = block_tables is not None and (

299 + block_tables.redirect_writes_to_null_block

300 + )

301 + if max_blocks_per_req > 0 and not null_blocks:

302 # Reserve block 0 (null block) and ensure we have enough blocks.

303 # Encoder-only models allocate no KV blocks, so this cap doesn't apply.

304 num_reqs = min(

@@ -450,4 +454,6 @@ def _warmup_kernels(

454 cleanup_output.finished_req_ids = set(req_ids)

455 worker_execute_model(cleanup_output)

456 model_runner.kv_connector.set_disabled(False)

457 + if model_runner.kv_block_zeroer is not None:

458 + model_runner.kv_block_zeroer.zero_block_ids([0])

459 torch.accelerator.synchronize()

vllm/v1/worker/gpu_model_runner.py

+17

@@ -4833,6 +4833,23 @@ class GPUModelRunner(

4833 self.input_batch.num_tokens_no_spec[i] = pos + 1

4834 self.input_batch.prev_req_id_to_index = prev_req_id_to_index

4835

4836 + @contextmanager

4837 + def preserve_serving_state(self):

4838 + multi_block_table = self.input_batch.block_table

4839 + saved = [

4840 + (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())

4841 + for bt in multi_block_table.block_tables

4842 + ]

4843 + multi_block_table.clear()

4844 + try:

4845 + yield

4846 + finally:

4847 + for bt, (saved_gpu, saved_cpu) in zip(

4848 + multi_block_table.block_tables, saved

4849 + ):

4850 + bt.block_table.gpu.copy_(saved_gpu)

4851 + bt.block_table.cpu.copy_(saved_cpu)

4852 +

4853 def take_draft_token_ids(self) -> DraftTokenIds | None:

4854 if not self.num_spec_tokens or not self._draft_token_req_ids:

4855 return None