[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>

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -224,7 +224,9 @@ def test_speculator_capture_preserves_decode_query_bounds(
     )
     manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)
     manager.max_num_reqs = num_reqs
-    manager.dp_size = 1
+    manager.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(data_parallel_size=1)
+    )
     buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
     block_tables = MagicMock()
     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
     client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]
     client.eng_start_index = 0
     client._kv_event_sources = {}
+    client.vllm_config = SimpleNamespace(
+        use_v2_model_runner=True,
+        parallel_config=SimpleNamespace(all2all_backend="allgather_reducescatter"),
+    )
     return client
 
 

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -24,7 +24,6 @@ class FakeMemoryProfiler:
 
 class FakeEplbState:
     instances: list["FakeEplbState"] = []
-    from_mapping_kwargs: dict[str, Any] | None = None
 
     def __init__(self, parallel_config: Any, device: torch.device):
         self.parallel_config = parallel_config
@@ -33,7 +32,7 @@ class FakeEplbState:
         self.step_calls: list[tuple[bool, bool, bool]] = []
         self.async_started = False
         self.is_async = True
-        self.built_from_mapping = False
+        self.update_mapping_calls: list[tuple[Any, torch.Tensor]] = []
         FakeEplbState.instances.append(self)
 
     def add_model(self, model: Any, model_config: Any) -> None:
@@ -45,12 +44,8 @@ class FakeEplbState:
     def start_async_loop(self) -> None:
         self.async_started = True
 
-    @classmethod
-    def from_mapping(cls, **kwargs: Any) -> "FakeEplbState":
-        cls.from_mapping_kwargs = kwargs
-        state = cls(kwargs["parallel_config"], kwargs["device"])
-        state.built_from_mapping = True
-        return state
+    def update_mapping(self, model_config: Any, mapping: torch.Tensor) -> None:
+        self.update_mapping_calls.append((model_config, mapping))
 
 
 def _make_runner(**overrides: Any) -> Any:
@@ -136,7 +131,7 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
     assert runner.eplb_state.async_started is True
 
 
-def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
+def test_v2_load_model_with_dummy_weights_defers_eplb_async_loop(monkeypatch):
     FakeEplbState.instances.clear()
     model = SimpleNamespace(is_moe=True)
 
@@ -159,24 +154,20 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
 
     assert runner.load_config.load_format == "dummy"
     assert runner.eplb_state is not None
-    assert runner.eplb_state.add_model_calls == []
+    # A scaling-up worker registers so it can join the EPLB communicator.
+    assert runner.eplb_state.add_model_calls == [(model, runner.model_config)]
     assert runner.eplb_state.async_started is False
 
 
-def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch):
-    FakeEplbState.instances.clear()
-    FakeEplbState.from_mapping_kwargs = None
-    monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
-    monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model)
-
-    runner = _make_runner(model=SimpleNamespace(is_moe=True))
+def test_v2_setup_eplb_from_mapping_updates_state_in_place():
+    runner = _make_runner()
+    state = FakeEplbState(runner.parallel_config, runner.device)
+    runner.eplb_state = state
     mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)
     mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping)
 
-    assert runner.eplb_state is not None
-    assert runner.eplb_state.built_from_mapping is True
-    assert FakeEplbState.from_mapping_kwargs is not None
-    assert FakeEplbState.from_mapping_kwargs["expanded_physical_to_logical"] is mapping
+    assert runner.eplb_state is state
+    assert state.update_mapping_calls == [(runner.model_config, mapping)]
 
 
 def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):

vllm/config/vllm.py

@@ -3086,9 +3086,6 @@ class VllmConfig:
         if self.parallel_config.use_ubatching:
             unsupported.extend(self._get_dbo_unsupported_features())
 
-        if self.parallel_config.enable_elastic_ep:
-            unsupported.append("elastic expert parallelism")
-
         if self.cache_config.mamba_cache_mode == "all":
             unsupported.append("mamba cache mode 'all'")
 

vllm/distributed/elastic_ep/elastic_execute.py

@@ -54,6 +54,8 @@ from vllm.platforms import current_platform
 from vllm.utils import is_moe_layer
 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
 from vllm.v1.worker.dp_utils import skip_dp_coordination
+from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
+from vllm.v1.worker.gpu.eplb_utils import preserve_serving_state
 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
 from vllm.v1.worker.workspace import lock_workspace, unlock_workspace
 
@@ -171,6 +173,17 @@ class ElasticEPScalingExecutor:
             self.worker.vllm_config.parallel_config,
         )
 
+    @contextmanager
+    def _suppress_eplb(self) -> Iterator[None]:
+        # A rank warming up alone must not advance the EPLB step counter.
+        runner = self.worker.model_runner
+        was_suppressed = runner.eep_eplb_suppressed
+        runner.eep_eplb_suppressed = True
+        try:
+            yield
+        finally:
+            runner.eep_eplb_suppressed = was_suppressed
+
     @contextmanager
     def _disable_flashinfer_autotune(self) -> Iterator[None]:
         kernel_config = self.worker.vllm_config.kernel_config
@@ -421,7 +434,17 @@ class ElasticEPScalingExecutor:
         self._staged_moe_quant_methods.clear()
 
     def _release_cuda_graphs(self) -> None:
-        if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
+        manager = getattr(self.worker.model_runner, "cudagraph_manager", None)
+        speculator = getattr(self.worker.model_runner, "speculator", None)
+        if manager is not None:
+            # MRV2 captures through CudaGraphManager instead of wrapping the
+            # model, so neither wrapper branch below ever fires.
+            manager.release_graphs()
+            for attr in vars(speculator).values() if speculator else ():
+                if isinstance(attr, CudaGraphManager):
+                    attr.release_graphs()
+
+        elif isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
             wrapper = self.worker.model_runner.model
             wrapper.concrete_cudagraph_entries = {}
 
@@ -431,6 +454,8 @@ class ElasticEPScalingExecutor:
         torch.compiler.reset()
         with set_current_vllm_config(self.worker.vllm_config):
             reset_compile_wrapper(self.worker.model_runner.get_model())
+            if speculator is not None:
+                reset_compile_wrapper(speculator.model)
 
         gc.collect()
         torch.accelerator.synchronize()
@@ -675,42 +700,28 @@ class ElasticEPScalingExecutor:
             self.warm_and_capture()
 
     def warm_and_capture(self) -> None:
-        # Save and clear block tables so the dummy MoE forward doesn't
-        # write dummy slot mappings into real KV-cache blocks.
-        multi_block_table = self.worker.model_runner.input_batch.block_table
-        saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
-        for bt in multi_block_table.block_tables:
-            saved_block_tables.append(
-                (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
-            )
-        multi_block_table.clear()
-
         # _ensure_workspace_size allocates a fresh tensor on grow, leaving
         # any captured CUDA graph with a stale data pointer; drop graphs
         # before re-warm so captures realign with the resized buffer.
         self._release_cuda_graphs()
         unlock_workspace()
 
-        # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model
-        # alone only exercises cudagraph-capture sizes and can leave the
-        # workspace too small for post-reshuffle routing. Use _dummy_run
-        # directly with skip_eplb=True so dummy routing doesn't pollute the
-        # just-rebalanced EPLB stats.
         runner = self.worker.model_runner
         all2all_manager = get_ep_all2all_manager()
         reuse_kernel = self._can_reuse_fused_moe_kernel()
+        serving_state = (
+            preserve_serving_state(runner)
+            if self.worker.use_v2_model_runner
+            else runner.preserve_serving_state()
+        )
         with (
             skip_dp_coordination() if reuse_kernel else nullcontext(),
             all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(),
             self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(),
+            self._suppress_eplb(),
+            serving_state,
         ):
             runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
             self.worker.compile_or_warm_up_model()
 
         lock_workspace()
-
-        for bt, (saved_gpu, saved_cpu) in zip(
-            multi_block_table.block_tables, saved_block_tables
-        ):
-            bt.block_table.gpu.copy_(saved_gpu)
-            bt.block_table.cpu.copy_(saved_cpu)

vllm/v1/engine/core_client.py

@@ -1893,13 +1893,29 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
         await asyncio.to_thread(self._coord_store.wait, ready_keys)
         logger.info("[Elastic EP] Successfully started new engines")
 
+    def _eep_commit_pause_mode(self) -> PauseMode:
+        from vllm.distributed.elastic_ep.elastic_execute import (
+            can_reuse_fused_moe_kernel,
+        )
+
+        # MRV2 re-warms through the request pool, so running requests must
+        # finish first while new ones stay queued in the scheduler.
+        parallel_config = self.vllm_config.parallel_config
+        if self.vllm_config.use_v2_model_runner and not can_reuse_fused_moe_kernel(
+            parallel_config
+        ):
+            return "wait"
+        return "keep"
+
     async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None:
         new_core_engines = [
             rank.to_bytes(2, "little")
             for rank in range(len(self.core_engines), new_data_parallel_size)
         ]
 
-        await self.pause_scheduler_async(mode="keep", clear_cache=False)
+        await self.pause_scheduler_async(
+            mode=self._eep_commit_pause_mode(), clear_cache=False
+        )
         wait_future = self._eep_wait_for_setup_switch_complete()
         finish_futures = [
             asyncio.create_task(
@@ -1971,7 +1987,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
         for rank in range(new_data_parallel_size, cur_data_parallel_size):
             self._kv_event_sources.pop(rank, None)
         removed_dp_size = cur_data_parallel_size - new_data_parallel_size
-        pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size
+        eep_mode = self._eep_commit_pause_mode()
+        pause_modes = [eep_mode] * new_data_parallel_size + ["abort"] * removed_dp_size
         pause_futures = [
             self._call_utility_async("pause_scheduler", mode, False, engine=engine)
             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 (
 
 
 class BlockTables:
+    # Set by the elastic EP warmup so KV writes land in the null block.
+    redirect_writes_to_null_block = False
+
     def __init__(
         self,
         block_sizes: list[int],
@@ -126,6 +129,8 @@ class BlockTables:
             bpk = self.blocks_per_kv_block[i]
             if bpk > 1:
                 block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
+            if self.redirect_writes_to_null_block:
+                block_ids = [0] * len(block_ids)
             end = start + len(block_ids)
             row_capacity = self.block_tables[i].gpu.shape[1]
             if end > row_capacity:

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -161,7 +161,6 @@ class CudaGraphManager:
         # DBO supports FULL CUDA graphs only.
         self.ubatch_runner = ubatch_runner
 
-        self.dp_size = vllm_config.parallel_config.data_parallel_size
         self.tp_size = vllm_config.parallel_config.tensor_parallel_size
         self.is_first_pp_rank = get_pp_group().is_first_rank
         self.is_last_pp_rank = get_pp_group().is_last_rank
@@ -385,6 +384,23 @@ class CudaGraphManager:
                         self._candidates.setdefault(key, []).extend(matching)
                     current_range_start = num_tokens + 1
 
+    @property
+    def dp_size(self) -> int:
+        # Not cached: elastic EP rewrites parallel_config in place on scale.
+        return self.vllm_config.parallel_config.data_parallel_size
+
+    def release_graphs(self) -> None:
+        """Drop the captured graphs so a later capture() can refill them.
+
+        Elastic EP reallocates the MoE workspace when it grows, which leaves
+        every captured graph holding a stale data pointer. `_capture_descs` is
+        kept, so `needs_capture()` still reports the work to redo.
+        """
+        self.graphs.clear()
+        self._graphs_captured = False
+        if self.breakable_cg_runner is not None:
+            BreakableCUDAGraphWrapper.clear_all_graphs()
+
     def needs_capture(self) -> bool:
         return len(self._capture_descs) > 0
 

vllm/v1/worker/gpu/eplb_utils.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
 from functools import wraps
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
@@ -15,6 +16,9 @@ from vllm.model_executor.models.interfaces import (
     get_mixture_of_experts_model,
 )
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
 logger = init_logger(__name__)
 
 
@@ -37,6 +41,27 @@ def step_eplb_after(*, is_dummy: bool = False) -> Callable:
     return decorator
 
 
+@contextmanager
+def preserve_serving_state(model_runner: "GPUModelRunner") -> Iterator[None]:
+    """Keep the elastic EP warmup out of the request pool and the KV cache."""
+    # After the drain only parked streaming sessions still hold a slot, and
+    # those are re-added from NewRequestData on their next chunk.
+    _remove_all_requests(model_runner)
+    model_runner.block_tables.redirect_writes_to_null_block = True
+    try:
+        yield
+    finally:
+        _remove_all_requests(model_runner)
+        model_runner.block_tables.redirect_writes_to_null_block = False
+        if model_runner.kv_block_zeroer is not None:
+            model_runner.kv_block_zeroer.zero_block_ids([0])
+
+
+def _remove_all_requests(model_runner: "GPUModelRunner") -> None:
+    for req_id in list(model_runner.req_states.req_id_to_index):
+        model_runner._remove_request(req_id)
+
+
 class EPLBController:
     def __init__(self, parallel_config: Any, device: torch.device):
         self.parallel_config = parallel_config
@@ -82,10 +107,8 @@ class EPLBController:
         self._has_registered_models = True
         return True
 
-    def maybe_register_model(
-        self, model: nn.Module, model_config: Any, load_dummy_weights: bool
-    ) -> bool:
-        if not self.parallel_config.enable_eplb or load_dummy_weights:
+    def maybe_register_model(self, model: nn.Module, model_config: ModelConfig) -> bool:
+        if not self.parallel_config.enable_eplb:
             return False
 
         moe_model = get_mixture_of_experts_model(model)
@@ -100,7 +123,13 @@ class EPLBController:
         self._has_registered_models = True
         return True
 
-    def maybe_start_async_loop(self, eplb_models_added: bool) -> None:
+    def maybe_start_async_loop(
+        self, eplb_models_added: bool, load_dummy_weights: bool = False
+    ) -> None:
+        # A scaling-up worker registers to join the communicator handshake, but
+        # its groups are not live yet, so _perform_eplb_reshuffle() starts this.
+        if load_dummy_weights:
+            return
         if eplb_models_added and self.state is not None and self.state.is_async:
             self.state.start_async_loop()
 
@@ -128,18 +157,10 @@ class EPLBController:
 
     def setup_from_mapping(
         self,
-        model: nn.Module,
-        model_config: Any,
+        model_config: ModelConfig,
         expanded_physical_to_logical: torch.Tensor,
     ) -> None:
-        moe_model = get_mixture_of_experts_model(model)
-        assert moe_model is not None
-
-        self.state = EplbState.from_mapping(
-            model=moe_model,
-            model_config=model_config,
-            device=self.device,
-            parallel_config=self.parallel_config,
-            expanded_physical_to_logical=expanded_physical_to_logical,
-        )
-        self._has_registered_models = True
+        # Update in place: a fresh EplbState would build a second communicator
+        # over a group the other ranks have already left.
+        assert self.state is not None
+        self.state.update_mapping(model_config, expanded_physical_to_logical)

vllm/v1/worker/gpu/model_runner.py

@@ -234,10 +234,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # Persistent buffer for intermediate tensors (non-first PP ranks).
         self.intermediate_tensors: IntermediateTensors | None = None
 
-        # Data parallelism.
-        self.dp_size = self.parallel_config.data_parallel_size
-        self.dp_rank = self.parallel_config.data_parallel_rank
-
         # Dual batch overlap. Created in initialize_kv_cache(), once everything
         # it runs the microbatched forward with exists.
         self.ubatch_runner: UBatchRunner | None = None
@@ -509,9 +505,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         if self.is_pooling_model and self.is_last_pp_rank:
             self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
         eplb_models_added |= self.eplb.maybe_register_model(
-            self.model, self.model_config, load_dummy_weights
+            self.model, self.model_config
         )
-        self.eplb.maybe_start_async_loop(eplb_models_added)
+        self.eplb.maybe_start_async_loop(eplb_models_added, load_dummy_weights)
 
         if not self.is_first_pp_rank:
             # For non-first PP ranks, create intermediate tensors sized
@@ -526,6 +522,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         get_offloader().post_init()
 
+    @property
+    def dp_size(self) -> int:
+        # Not cached: elastic EP rewrites parallel_config in place on reconfigure.
+        return self.parallel_config.data_parallel_size
+
+    @property
+    def dp_rank(self) -> int:
+        # Not cached: elastic EP rewrites parallel_config in place on reconfigure.
+        return self.parallel_config.data_parallel_rank
+
     def get_model(self) -> nn.Module:
         return self.model
 
@@ -2337,9 +2343,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def setup_eplb_from_mapping(
         self, expanded_physical_to_logical: torch.Tensor
     ) -> None:
-        self.eplb.setup_from_mapping(
-            self.model, self.model_config, expanded_physical_to_logical
-        )
+        self.eplb.setup_from_mapping(self.model_config, expanded_physical_to_logical)
 
     ########### EPLB methods end ###########
 

vllm/v1/worker/gpu/warmup.py

@@ -294,7 +294,11 @@ def _warmup_kernels(
         model_runner.scheduler_config.max_num_batched_tokens
         // max(prompt_len, decode_query_len),
     )
-    if max_blocks_per_req > 0:
+    block_tables = getattr(model_runner, "block_tables", None)
+    null_blocks = block_tables is not None and (
+        block_tables.redirect_writes_to_null_block
+    )
+    if max_blocks_per_req > 0 and not null_blocks:
         # Reserve block 0 (null block) and ensure we have enough blocks.
         # Encoder-only models allocate no KV blocks, so this cap doesn't apply.
         num_reqs = min(
@@ -450,4 +454,6 @@ def _warmup_kernels(
     cleanup_output.finished_req_ids = set(req_ids)
     worker_execute_model(cleanup_output)
     model_runner.kv_connector.set_disabled(False)
+    if model_runner.kv_block_zeroer is not None:
+        model_runner.kv_block_zeroer.zero_block_ids([0])
     torch.accelerator.synchronize()

vllm/v1/worker/gpu_model_runner.py

@@ -4833,6 +4833,23 @@ class GPUModelRunner(
             self.input_batch.num_tokens_no_spec[i] = pos + 1
         self.input_batch.prev_req_id_to_index = prev_req_id_to_index
 
+    @contextmanager
+    def preserve_serving_state(self):
+        multi_block_table = self.input_batch.block_table
+        saved = [
+            (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
+            for bt in multi_block_table.block_tables
+        ]
+        multi_block_table.clear()
+        try:
+            yield
+        finally:
+            for bt, (saved_gpu, saved_cpu) in zip(
+                multi_block_table.block_tables, saved
+            ):
+                bt.block_table.gpu.copy_(saved_gpu)
+                bt.block_table.cpu.copy_(saved_cpu)
+
     def take_draft_token_ids(self) -> DraftTokenIds | None:
         if not self.num_spec_tokens or not self._draft_token_req_ids:
             return None