[Elastic EP] Support Model Runner V2

tests/v1/engine/test_admission_control.py

@@ -382,3 +382,101 @@ def test_human_readable_int_parses_notation(input_str: str, expected: int):
 def test_human_readable_int_rejects_invalid(invalid: str):
     with pytest.raises((argparse.ArgumentTypeError, ValueError)):
         human_readable_int(invalid)
+
+
+@pytest.mark.asyncio
+async def test_add_request_rechecks_admission_after_middleware_passes():
+    llm = _make_async_llm()
+    llm.output_processor.has_request.return_value = False
+    llm.engine_core = SimpleNamespace(
+        add_request_async=AsyncMock(), shutdown=MagicMock()
+    )
+    request = SimpleNamespace(request_id="late")
+
+    from vllm.entrypoints.serve.elastic_ep.middleware import (
+        set_scaling_elastic_ep,
+    )
+
+    set_scaling_elastic_ep(False)
+    llm.check_admission(request_id=request.request_id)
+    set_scaling_elastic_ep(True)
+    try:
+        with pytest.raises(GracefulHTTPError) as exc_info:
+            await llm._add_request(request, None, None, 0, MagicMock())
+        assert exc_info.value.http_status == HTTPStatus.SERVICE_UNAVAILABLE
+        llm.engine_core.add_request_async.assert_not_awaited()
+    finally:
+        set_scaling_elastic_ep(False)
+
+
+@pytest.mark.asyncio
+async def test_existing_request_can_add_after_admission_closes():
+    llm = _make_async_llm()
+    llm.output_processor.has_request.return_value = True
+    llm.engine_core = SimpleNamespace(
+        add_request_async=AsyncMock(), shutdown=MagicMock()
+    )
+    llm.log_requests = False
+    request = SimpleNamespace(request_id="streaming-continuation")
+
+    from vllm.entrypoints.serve.elastic_ep.middleware import (
+        set_scaling_elastic_ep,
+    )
+
+    set_scaling_elastic_ep(True)
+    try:
+        await llm._add_request(request, None, None, 0, MagicMock())
+        llm.engine_core.add_request_async.assert_awaited_once_with(request)
+    finally:
+        set_scaling_elastic_ep(False)
+
+
+@pytest.mark.asyncio
+async def test_parallel_sampling_registers_all_children_before_first_send(
+    monkeypatch,
+):
+    """A drain that starts mid fan-out must see every child of the parent."""
+    llm = _make_async_llm()
+    llm.output_handler = None
+    llm.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(kv_sharing_fast_prefill=False)
+    )
+    llm.model_config = MagicMock()
+    llm.log_requests = False
+    llm._run_output_handler = MagicMock()
+    llm.get_supported_tasks = AsyncMock(return_value=())
+    params = SamplingParams(n=3)
+    request = SimpleNamespace(
+        request_id="parent",
+        external_req_id="parent",
+        params=params,
+        sampling_params=params,
+    )
+    llm.input_processor = MagicMock()
+    llm.input_processor.process_inputs_async = AsyncMock(return_value=request)
+    monkeypatch.setattr(
+        "vllm.v1.engine.async_llm.extract_prompt_components",
+        lambda *_: ("hi", None, None),
+    )
+    first_send_started = asyncio.Event()
+    release_first_send = asyncio.Event()
+
+    async def add_request_async(_request):
+        first_send_started.set()
+        await release_first_send.wait()
+
+    llm.engine_core = SimpleNamespace(
+        resources=SimpleNamespace(engine_dead=False),
+        add_request_async=AsyncMock(side_effect=add_request_async),
+        shutdown=MagicMock(),
+    )
+
+    add_task = asyncio.create_task(llm.add_request("parent", "hi", params))
+    try:
+        await first_send_started.wait()
+        assert llm.output_processor.add_request.call_count == 3
+        assert llm.engine_core.add_request_async.await_count == 1
+    finally:
+        release_first_send.set()
+        await add_task
+    assert llm.engine_core.add_request_async.await_count == 3

tests/v1/worker/test_gpu_warmup_blocks.py

@@ -107,6 +107,8 @@ def _make_runner(
         ),
         kv_block_zeroer=None,
         kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
+        req_states=SimpleNamespace(free_indices=list(range(4))),
+        block_tables=None,
     )
 
 

vllm/config/vllm.py

@@ -2790,9 +2790,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
         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")
-
         has_logitsproc_plugins = False
         if model_config is not None:
             from importlib.metadata import entry_points

vllm/distributed/elastic_ep/elastic_execute.py

@@ -170,6 +170,17 @@ def _can_reuse_fused_moe_kernel(self) -> bool:
             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
@@ -418,7 +429,13 @@ def _commit_staged_moe_quant_methods(self) -> None:
         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)
+        if manager is not None:
+            # MRV2 captures through CudaGraphManager instead of wrapping the
+            # model, so neither wrapper branch below ever fires.
+            manager.release_graphs()
+
+        elif isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
             wrapper = self.worker.model_runner.model
             wrapper.concrete_cudagraph_entries = {}
 
@@ -667,42 +684,23 @@ def warmup_new_worker(self) -> None:
             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()
         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(),
+            runner.preserve_serving_state(),
         ):
-            runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
+            runner.warm_up_workspace()
             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/async_llm.py

@@ -7,6 +7,7 @@
 import warnings
 from collections.abc import AsyncGenerator, Iterable, Mapping
 from copy import copy
+from http import HTTPStatus
 from typing import Any
 
 import vllm.envs as envs
@@ -18,7 +19,10 @@
 )
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.protocol import EngineClient, StreamingInput
-from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep
+from vllm.entrypoints.serve.elastic_ep.middleware import (
+    get_scaling_elastic_ep,
+    set_scaling_elastic_ep,
+)
 from vllm.exceptions import (
     GracefulHTTPError,
     MaxQueuedTokensError,
@@ -310,6 +314,8 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None:
             QueueOverflowError: If ``max_num_queued_reqs`` would be exceeded.
             MaxQueuedTokensError: If ``max_num_queued_tokens`` would be exceeded.
         """
+        self._check_elastic_ep_admission(request_id)
+
         max_num_reqs = self.scheduler_config.max_num_queued_reqs
         if max_num_reqs is not None:
             current = self.get_num_unfinished_requests()
@@ -337,6 +343,17 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None:
                 )
                 raise MaxQueuedTokensError()
 
+    def _check_elastic_ep_admission(self, request_id: str | None = None) -> None:
+        if get_scaling_elastic_ep():
+            logger.info(
+                "Elastic EP scaling is in progress - rejecting request %s",
+                request_id,
+            )
+            raise GracefulHTTPError(
+                "The model is currently scaling. Please try again later.",
+                HTTPStatus.SERVICE_UNAVAILABLE,
+            )
+
     async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         if not hasattr(self, "_supported_tasks"):
             # Cache the result
@@ -480,14 +497,22 @@ async def add_request(
 
         # Fan out child requests (for n>1).
         parent_request = ParentRequest(request)
+        self._check_elastic_ep_admission(request.request_id)
+        child_requests = []
         for idx in range(parent_params.n):
             request_id, child_params = parent_request.get_child_info(idx)
             child_request = request if idx == parent_params.n - 1 else copy(request)
             child_request.request_id = request_id
             child_request.sampling_params = child_params
-            await self._add_request(
+            # Register every child before the first await so a drain sees all of them.
+            self.output_processor.add_request(
                 child_request, prompt_text, parent_request, idx, queue
             )
+            child_requests.append(child_request)
+        for child_request in child_requests:
+            await self.engine_core.add_request_async(child_request)
+            if self.log_requests:
+                logger.info("Added request %s.", child_request.request_id)
         return queue
 
     async def _add_request(
@@ -1113,11 +1138,18 @@ async def wait_for_requests_to_drain(self, drain_timeout: int = 300):
         """Wait for all requests to be drained."""
         start_time = time.time()
         while time.time() - start_time < drain_timeout:
-            if not self.engine_core.dp_engines_running():
+            dp_engines_running = self.engine_core.dp_engines_running()
+            has_unfinished_requests = self.output_processor.has_unfinished_requests()
+            if not dp_engines_running and not has_unfinished_requests:
                 logger.info("Engines are idle, requests have been drained")
                 return
 
-            logger.info("Engines are still running, waiting for requests to drain...")
+            logger.info(
+                "Waiting for requests to drain "
+                "(engines_running=%s, frontend_unfinished=%s)",
+                dp_engines_running,
+                has_unfinished_requests,
+            )
             await asyncio.sleep(1)  # Wait 1 second before checking again
 
         raise TimeoutError(
@@ -1127,10 +1159,7 @@ async def wait_for_requests_to_drain(self, drain_timeout: int = 300):
 
     async def _drain_requests_for_elastic_ep(self, drain_timeout: int) -> None:
         try:
-            logger.info(
-                "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, "
-                "waiting for requests to drain before scaling"
-            )
+            logger.info("Waiting for requests to drain before Elastic EP scaling")
             await self.wait_for_requests_to_drain(drain_timeout)
         except BaseException:
             set_scaling_elastic_ep(False)
@@ -1173,8 +1202,17 @@ async def _scale_elastic_ep(
                 self._logger_ref[0] = self.logger_manager
             self.logger_manager.log_engine_initialized()
 
+        from vllm.distributed.elastic_ep.elastic_execute import (
+            can_reuse_fused_moe_kernel,
+        )
+
+        # MRV2 ranks that re-warm at commit need an empty request pool.
+        drain = envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS or (
+            self.vllm_config.use_v2_model_runner
+            and not can_reuse_fused_moe_kernel(self.vllm_config.parallel_config)
+        )
         set_scaling_elastic_ep(True)
-        if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS:
+        if drain:
             await self._drain_requests_for_elastic_ep(drain_timeout)
 
         await self.engine_core.commit_elastic_ep()

vllm/v1/worker/gpu/block_table.py

@@ -82,6 +82,9 @@ def __init__(
             device=self.device,
         )
 
+        # Set by the elastic EP warmup so KV writes land in the null block.
+        self.redirect_writes_to_null_block = False
+
         self.init_block_table_layout_tensors()
 
     def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
@@ -122,6 +125,8 @@ def append_block_ids(
             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

@@ -154,7 +154,6 @@ def __init__(
         self.decode_query_len = decode_query_len
         self.varlen_decode = varlen_decode
 
-        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
@@ -341,6 +340,23 @@ def _init_candidates(self) -> None:
                         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