@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 17da48596c

[Bugfix][DP] Synchronize the device on pause completion (#52914)

Signed-off-by: aoshen02 <aoshen@inferact.ai>

aoshen02 committed Aug 27, 2026 at 00:50 UTC 17da48596c98946d3e3e6896e2ebd341e809f3bd

5 files changed +164 -8

tests/v1/distributed/test_async_llm_dp.py

+110

@@ -713,3 +713,113 @@ async def test_dp_pause_barrier_request_deadlock():

713 assert not await engine.is_paused()

714 # Let the two requests we sent mid-barrier complete.

715 await asyncio.gather(*mid_barrier_tasks)

716 +

717 +

718 +@pytest.mark.asyncio

719 +async def test_dp_pause_wait_mode_drains_in_flight():

720 + """mode="wait" through the DP consensus path: the pause lets the

721 + in-flight request run to completion while the consensus is pending,

722 + resolves only once drained, and the engine generates again on resume."""

723 + with ExitStack() as after:

724 + engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True))

725 + after.callback(engine.shutdown)

726 +

727 + collector = await engine.add_request(

728 + request_id="inflight",

729 + prompt=DP_PAUSE_PROMPT,

730 + params=SamplingParams(max_tokens=64),

731 + )

732 + await engine.pause_generation(mode="wait")

733 + assert await engine.is_paused()

734 + while True:

735 + out = await asyncio.wait_for(collector.get(), timeout=30)

736 + if out.finished:

737 + break

738 +

739 + await engine.resume_generation()

740 + async for out in engine.generate(

741 + request_id="after-wait",

742 + prompt=DP_PAUSE_PROMPT,

743 + sampling_params=SamplingParams(max_tokens=5),

744 + ):

745 + pass

746 + assert out.finished

747 +

748 +

749 +@pytest.mark.asyncio

750 +async def test_dp_pause_while_asleep():

751 + """Pausing a sleeping DP engine: sleeping ranks skip dummy batches yet

752 + must still reach the pause consensus, and the completion barrier must

753 + be harmless on workers whose memory is unmapped."""

754 + with ExitStack() as after:

755 + engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True))

756 + after.callback(engine.shutdown)

757 +

758 + async for _ in engine.generate(

759 + request_id="warmup",

760 + prompt=DP_PAUSE_PROMPT,

761 + sampling_params=SamplingParams(max_tokens=5),

762 + ):

763 + pass

764 + await engine.sleep(level=1)

765 + assert await engine.is_sleeping()

766 +

767 + await asyncio.wait_for(engine.pause_generation(mode="abort"), timeout=30)

768 + assert await engine.is_paused()

769 +

770 + # A full wake makes the memory resident and resumes the scheduler.

771 + await engine.wake_up()

772 + assert not await engine.is_sleeping()

773 + async for out in engine.generate(

774 + request_id="after-wake",

775 + prompt=DP_PAUSE_PROMPT,

776 + sampling_params=SamplingParams(max_tokens=5),

777 + ):

778 + pass

779 + assert out.finished

780 +

781 +

782 +@pytest.mark.asyncio

783 +async def test_dp_pause_completion_implies_device_idle(

784 + monkeypatch: pytest.MonkeyPatch,

785 +):

786 + """A resolved pause future promises an idle device. Inflate the dummy

787 + batches the pause consensus manufactures so any work the pause fails to

788 + wait on stays visible, then require a quiet worker stream the moment

789 + pause_generation() returns."""

790 + with ExitStack() as after:

791 + # The probes below ship functions through collective_rpc.

792 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

793 + engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True))

794 + after.callback(engine.shutdown)

795 +

796 + async for _ in engine.generate(

797 + request_id="warmup",

798 + prompt=DP_PAUSE_PROMPT,

799 + sampling_params=SamplingParams(max_tokens=5),

800 + ):

801 + pass

802 +

803 + def inflate_dummy_batches(worker: Any) -> bool:

804 + import torch

805 +

806 + inner = worker.execute_dummy_batch

807 +

808 + def slow_dummy_batch() -> None:

809 + inner()

810 + # Keep the stream busy after launch, like a large-MoE forward.

811 + torch.cuda._sleep(200_000_000)

812 +

813 + worker.execute_dummy_batch = slow_dummy_batch

814 + return True

815 +

816 + assert all(await engine.engine_core.collective_rpc_async(inflate_dummy_batches))

817 +

818 + await asyncio.wait_for(engine.pause_generation(mode="abort"), timeout=60)

819 +

820 + def device_idle(worker: Any) -> bool:

821 + import torch

822 +

823 + return bool(torch.cuda.current_stream().query())

824 +

825 + assert all(await engine.engine_core.collective_rpc_async(device_idle))

tests/v1/engine/test_engine_core.py

+36 -2

@@ -5,7 +5,7 @@ import copy

5 import time

6 import uuid

7 from concurrent.futures import Future, ThreadPoolExecutor

8 -from unittest.mock import PropertyMock, patch

8 +from unittest.mock import MagicMock, PropertyMock, patch

10 import pytest

11 from transformers import AutoTokenizer

@@ -23,7 +23,7 @@ from vllm.engine.arg_utils import EngineArgs

23 from vllm.platforms import current_platform

24 from vllm.utils.torch_utils import set_default_torch_num_threads

25 from vllm.v1.engine import EngineCoreRequest

26 -from vllm.v1.engine.core import EngineCore

26 +from vllm.v1.engine.core import EngineCore, EngineCoreProc

27 from vllm.v1.executor.abstract import Executor

28 from vllm.v1.executor.uniproc_executor import UniProcExecutor

29 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -608,3 +608,37 @@ def test_encoder_instance_zero_kv_cache(

608 assert not engine_core.scheduler.ec_connector.is_producer, (

609 "Consumer instance EC connector should be consumer"

610 )

611 +

612 +

613 +def _pausable_engine_core_proc() -> EngineCoreProc:

614 + """A bare EngineCoreProc holding just the state pause_scheduler touches."""

615 + core = object.__new__(EngineCoreProc)

616 + core.model_executor = MagicMock()

617 + core.scheduler = MagicMock()

618 + core.scheduler.has_requests.return_value = False

619 + core.batch_queue = None

620 + core.engines_running = False

621 + core._idle_state_callbacks = []

622 + return core

623 +

624 +

625 +@pytest.mark.parametrize("deferred", [False, True])

626 +def test_pause_synchronizes_device_before_cache_reset(deferred: bool):

627 + """A resolved pause promises an idle device: the barrier must run before

628 + caches are cleared and before the caller is unblocked."""

629 + core = _pausable_engine_core_proc()

630 + core.engines_running = deferred

631 + order: list[str] = []

632 + core.model_executor.collective_rpc.side_effect = lambda method: order.append(method)

633 + core._reset_caches = lambda: order.append("reset_caches")

634 +

635 + result = EngineCoreProc.pause_scheduler(core, mode="keep", clear_cache=True)

636 + if deferred:

637 + assert isinstance(result, Future)

638 + assert not result.done() and order == []

639 + core.engines_running = False

640 + core._notify_idle_state_callbacks()

641 + assert result.result(timeout=0) is None

642 + else:

643 + assert result is None

644 + assert order == ["synchronize_device", "reset_caches"]

vllm/v1/engine/core.py

+10 -6

@@ -839,6 +839,13 @@ class EngineCore:

839 self.reset_mm_cache()

840 self.reset_encoder_cache()

841

842 + def _finish_pause(self, clear_cache: bool) -> None:

843 + # A completed pause promises an idle device: nothing else waits on

844 + # the last dummy batch an idle DP rank launches.

845 + self.model_executor.collective_rpc("synchronize_device")

846 + if clear_cache:

847 + self._reset_caches()

848 +

849 def pause_scheduler(

850 self, mode: PauseMode = "abort", clear_cache: bool = True

851 ) -> Future | None:

@@ -865,8 +872,7 @@ class EngineCore:

872

873 pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW

874 self.scheduler.set_pause_state(pause_state)

868 - if clear_cache:

869 - self._reset_caches()

875 + self._finish_pause(clear_cache)

876

877 return None

878

@@ -1937,8 +1943,7 @@ class EngineCoreProc(EngineCore):

1943 raise ValueError(f"Invalid pause mode: {mode}")

1944

1945 def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:

1940 - if clear_cache:

1941 - engine._reset_caches()

1946 + engine._finish_pause(clear_cache)

1947 future.set_result(None)

1948

1949 if mode == "abort":

@@ -1951,8 +1956,7 @@ class EngineCoreProc(EngineCore):

1956 self.scheduler.set_pause_state(pause_state)

1957

1958 if self._pause_complete():

1954 - if clear_cache:

1955 - self._reset_caches()

1959 + self._finish_pause(clear_cache)

1960 return None

1961

1962 future = Future[Any]()

vllm/v1/worker/gpu_worker.py

@@ -259,6 +259,8 @@ class Worker(WorkerBase):

259 if tags is None or "kv_cache" in tags:

260 self.model_runner.post_kv_cache_wake_up()

261

262 + self.synchronize_device()

263 +

264 def checkpoint_prepare(self) -> None:

265 checkpoint_prepare_distributed_state()

266

vllm/v1/worker/worker_base.py

@@ -125,6 +125,12 @@ class WorkerBase:

125 """Basic health check (override for device-specific checks)."""

126 return

127

128 + def synchronize_device(self) -> None:

129 + """Block until in-flight device work completes; backends outside

130 + ``torch.accelerator`` must override with their own wait."""

131 + if torch.accelerator.is_available():

132 + torch.accelerator.synchronize()

133 +

134 def init_device(self) -> None:

135 """Initialize device state, such as loading the model or other on-device

136 memory allocations.