[Core][Diffusion] Add pause_generation(mode="keep") for diffusion stages (#7685)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

docs/features/sleep_mode.md

@@ -9,9 +9,11 @@ This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/1
 ## 1. Feature Documentation
 
 ### Overview
+
 Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization.
 
 ### Sleep Levels
+
 We support two levels of hibernation to balance recovery speed and memory efficiency:
 
 | Level | Name | Mechanism | Recovery Speed | Memory Freed |
@@ -29,15 +31,67 @@ Omni Sleep Mode is optimized for high-performance computing backends:
 * **Huawei NPU**: Supported via Ascend memory scavenging.
 
 ### Hardware Requirements
+
 * **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep.
 * **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers.
 
----
+### Quiescing a busy diffusion stage: `pause_generation(mode="keep")`
 
+`sleep()` on a diffusion stage is a worker-level RPC: it offloads memory but does not stop the
+stage's scheduler, so a stage that still has queued requests must be made quiet first. Aborting
+discards the batch that is running. `pause_generation(mode="keep")` keeps it instead:
+
+* The batch that is already running finishes on the weights it started with and its outputs are
+  delivered normally.
+* Requests that have not started stay queued and run after `resume_generation()`, on whatever
+  weights are loaded by then. A request never continues denoising across a weight change.
+* The call returns only after every replica has acknowledged: the batch that was running has
+  finished executing, its asynchronous output copies are done, and the device is synchronized; its
+  outputs are then delivered on the normal path. On a non-preemptible backend that acknowledgement
+  arrives at the next batch boundary, so expect it to take as long as the remaining batch.
+* Control RPCs (`list_loras`, `sleep`, `wake_up`, `resume_generation`) keep working while paused.
+
+Recommended order for a weight update on a diffusion stage:
+
+```python
+await engine.pause_generation(mode="keep", clear_cache=False)
+await engine.sleep(level=1)          # optional: free memory for the trainer
+await engine.wake_up()               # admission stays closed until resume
+# install new weights here
+await engine.resume_generation()
+```
+
+Support matrix and limits:
+
+| Configuration | `mode="keep"` |
+| :--- | :--- |
+| Request-level execution (`step_execution=False`) with the in-process (`uni`) executor | Supported |
+| Request-level execution with the multi-process (`mp`) executor, including asynchronous output | Supported; the acknowledgement waits for the output copies |
+| Request-level execution with asynchronous KV prefetch | Supported; the acknowledgement joins any prefetch still in flight |
+| Step-level execution (`step_execution=True` / streaming output) | Not supported: the call raises `NotImplementedError` |
+| `mode="abort"` / `mode="wait"` / `wait_for_inflight_requests=True` on a diffusion stage | Unchanged: frontend admission is paused only |
+
+Behavior notes:
+
+* Callers that already passed `mode="keep"` for a diffusion stage used to get an immediate return
+  that only closed frontend admission. They now also stop the backend and wait for the ACK.
+* A timeout or error from `pause_generation` means the pause did **not** complete; the scheduler
+  stays closed and the caller must retry the pause or call `resume_generation()` explicitly. Do not
+  sleep or mutate weights after a failed pause.
+* `resume_generation(stage_ids=...)` that leaves another paused stage closed (an AR stage, or a
+  keep-paused diffusion stage) keeps frontend admission closed as well, so new requests cannot
+  queue on a stage that will not schedule them. Admission reopens once every stage paused by
+  `pause_generation` has been resumed.
+* `clear_cache` has no diffusion-specific effect.
+* Without a pause call nothing changes: no barrier runs and batching, abort and output ordering are
+  as before.
+
+---
 
 ## 2. Usage Examples
 
 ### Python API Example
+
 You can programmatically control the lifecycle of stages using the `AsyncOmni` engine.
 
 ```python
@@ -65,6 +119,7 @@ if __name__ == "__main__":
 ```
 
 ### server command Example
+
 Start the server with sleep mode enabled:
 
 The first method
@@ -92,9 +147,6 @@ python3 -m vllm_omni.entrypoints.openai.api_server \
 
 ```
 
-
-
-
 ### Test Scenarios & Commands
 
 #### Scenario 1: LLM Engine Sleep
@@ -113,9 +165,8 @@ curl -X POST http://localhost:8000/v1/omni/sleep \
 
 Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage.
 
-
-
 #### Scenario 2: Diffusion Sleep
+
 Objective: Verify VRAM reclamation for Stage 1 (Diffusion).
 
 Trigger sleep (Level 1 or Level 2) via client:
@@ -128,9 +179,8 @@ curl -X POST http://localhost:8000/v1/omni/sleep \
 
 ```
 
-
-
 #### Scenario 3: Multi-Stage Coordinated Stress Test
+
 Objective: Test concurrent sleep and rapid wake-up across multiple stages.
 
 Concurrent Sleep (Stage 0 & 1):
@@ -143,7 +193,6 @@ curl -X POST http://localhost:8000/v1/omni/sleep \
 
 ```
 
-
 Rapid Wake-up:
 
 ```
@@ -154,8 +203,8 @@ curl -X POST http://localhost:8000/v1/omni/wakeup \
 
 ```
 
-
 #### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity
+
 Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation.
 
 Check Initial State: Observe baseline VRAM usage.
@@ -196,20 +245,16 @@ curl -X POST http://localhost:8000/v1/images/generations \
 
 ```
 
-
-
-
 ## 3. API Reference
 
-
 ### Methods
 
 | Method | Arguments | Return Type | Description |
 | :--- | :--- | :--- | :--- |
 | **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. |
 | **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. |
-
-
+| **pause_generation** | `mode: str, stage_ids: List[int]` | `None` | Stops admission; with `mode="keep"` also stops diffusion schedulers and returns after their ACK. |
+| **resume_generation** | `stage_ids: List[int]` | `None` | Reopens paused schedulers, then admission. |
 
 ### OmniACK Dataclass Fields
 
@@ -233,7 +278,8 @@ The metadata field is a dynamic dictionary containing hardware-specific telemetr
 }
 ```
 
-#### Core Utility:
+#### Core Utility
+
 **VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM.
 
 **Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines.

tests/diffusion/distributed/test_pause_keep_two_ranks.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Two-rank ``pause_generation(mode="keep")``: every rank runs the pause barrier
+before the ACK, and a barrier failure on a non-zero rank fails the pause
+without aborting or dropping the running batch.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from contextlib import ExitStack
+
+import pytest
+
+from tests.e2e.features.custom_pipeline.test_diffusion_pause_keep import (
+    CUSTOM_PIPELINE_CLASS,
+    MODEL,
+    WORKER_EXTENSION_CLASS,
+    _generate,
+    _pause_round,
+    _wait_for,
+)
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+
+def _all_true(results) -> bool:
+    flat = []
+    stack = list(results)
+    while stack:
+        item = stack.pop()
+        if isinstance(item, list):
+            stack.extend(item)
+        else:
+            flat.append(item)
+    return bool(flat) and all(item is True for item in flat)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_keep_pause_two_ranks_ack_and_rank_failure():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            num_gpus=2,
+            parallel_config={"ulysses_degree": 2, "ulysses_mode": "advanced_uaa"},
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+        assert _all_true(await engine.collective_rpc(method="start_counting_synchronize_device"))
+
+        pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(3)]
+        # The result is the all-rank AND, so True means both ranks ran the barrier.
+        assert _all_true(await engine.collective_rpc(method="synchronize_device_seen"))
+
+        assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, True)))
+        a = asyncio.create_task(_generate(engine, "a-fail", events))
+        await _wait_for(lambda: "started" in events["a-fail"])
+        with pytest.raises(RuntimeError, match="pause_scheduler failed"):
+            await engine.pause_generation(mode="keep", clear_cache=False)
+        assert await engine.is_paused()
+        a_output = await asyncio.wait_for(a, 60.0)
+        assert a_output.images, "a failed barrier must not abort the running batch"
+
+        assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, False)))
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+        recovered = await asyncio.wait_for(_generate(engine, "after-recovery", events), 60.0)
+        assert recovered.images
+
+        assert not engine.request_states
+        print(f"[sp2] rounds=3 pause_latency_s min/max={min(pause_latencies):.3f}/{max(pause_latencies):.3f}")

tests/diffusion/test_async_output_worker.py

@@ -497,3 +497,99 @@ class TestWorkerProcShutdown:
         result_mq.shutdown.assert_called_once_with()
         assert proc.mq is None
         assert proc.result_mq is None
+
+
+class TestSynchronizeDeviceBarrier:
+    """The pause barrier must not report a quiet device while outputs still pack."""
+
+    @staticmethod
+    def _rpc(timeout):
+        return {
+            "method": "synchronize_device",
+            "args": (),
+            "kwargs": {"timeout": timeout},
+            "output_rank": 0,
+            "exec_all_ranks": False,
+            "collect_rank_status": False,
+        }
+
+    def test_rejects_undrained_output(self, mocker):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        mocker.patch.object(proc, "drain_async_outputs", return_value=False)
+
+        with pytest.raises(TimeoutError, match="did not drain"):
+            proc._execute_rpc(self._rpc(0.01))
+
+        proc.worker.execute_method.assert_not_called()
+
+    @pytest.mark.parametrize("timeout", [0.5, None])
+    def test_drains_with_the_call_timeout_then_synchronizes(self, mocker, timeout):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        call_order: list = []
+        mocker.patch.object(
+            proc,
+            "drain_async_outputs",
+            side_effect=lambda **kw: call_order.append(("drain", kw)) or True,
+        )
+        proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append(("sync", a, kw))
+
+        _, should_reply = proc._execute_rpc(self._rpc(timeout))
+
+        assert call_order == [
+            ("drain", {"timeout": timeout}),
+            ("sync", ("synchronize_device",), {"timeout": timeout}),
+        ]
+        assert should_reply is True
+
+    def test_completes_only_after_output_ready_is_published(self, mocker):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+
+        mock_platform = mocker.MagicMock()
+        mock_platform.record_device_event.return_value = MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
+            mock_platform,
+        )
+        mock_torch = mocker.MagicMock()
+        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
+        release = threading.Event()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=lambda *a, **kw: release.wait(timeout=5.0),
+        )
+
+        proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
+        loop_thread = threading.Thread(target=proc._async_output_loop, daemon=True)
+        loop_thread.start()
+
+        barrier_done = threading.Event()
+
+        def run_barrier():
+            proc._execute_rpc(self._rpc(5.0))
+            barrier_done.set()
+
+        barrier_thread = threading.Thread(target=run_barrier, daemon=True)
+        barrier_thread.start()
+        try:
+            assert not barrier_done.wait(timeout=0.2)
+            proc.worker.execute_method.assert_not_called()
+
+            release.set()
+            assert barrier_done.wait(timeout=5.0)
+
+            kinds = [
+                call.args[0].kind
+                for call in proc.result_mq.enqueue.call_args_list
+                if isinstance(call.args[0], AsyncDiffusionOutput)
+            ]
+            assert kinds == [AsyncOutputKind.COMPUTE_DONE, AsyncOutputKind.OUTPUT_READY]
+            proc.worker.execute_method.assert_called_once_with("synchronize_device", timeout=5.0)
+        finally:
+            release.set()
+            proc._async_output_queue.put(None)
+            loop_thread.join(timeout=2.0)
+            barrier_thread.join(timeout=2.0)

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for the busy-loop-based RPC routing in DiffusionEngine.
 
 These tests guard the invariants established when ``collective_rpc`` was
@@ -60,6 +60,12 @@ class _ConcurrencyTrackingExecutor:
         self.is_failed = False
         self._closed = False
         self.od_config = SimpleNamespace()
+        # Pause tests: a call whose gate is registered (by request id, or by
+        # method name for ``synchronize_device``) blocks inside the executor
+        # until the gate is set; ``version`` tags each call so a test can tell
+        # which "weights" a request ran on.
+        self.gates: dict[str, threading.Event] = {}
+        self.version = "v1"
 
     def collective_rpc(
         self,
@@ -81,16 +87,25 @@ class _ConcurrencyTrackingExecutor:
                     "kwargs": kwargs,
                     "unique_reply_rank": unique_reply_rank,
                     "thread": threading.get_ident(),
+                    "version": self.version,
                 }
             )
         try:
             if self.rpc_delay:
                 time.sleep(self.rpc_delay)
+            if method == "synchronize_device":
+                gate = self.gates.get(method)
+                if gate is not None:
+                    gate.wait(timeout=5.0)
+                return [None]
             # Distinguish per-request execution from raw RPC by method name.
             if method in {"execute_model", "execute_stepwise", "generate"}:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
                 tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
+                gate = self.gates.get(tag)
+                if gate is not None:
+                    gate.wait(timeout=5.0)
                 return DiffusionOutput(error=f"result_for_{tag}", finished=True)
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -133,21 +148,27 @@ def _make_request(tag: str):
 
 
 def _make_engine_with_loop(
-    loop: asyncio.AbstractEventLoop,
+    loop: asyncio.AbstractEventLoop | None,
     rpc_delay: float = 0.0,
+    *,
+    scheduler: RequestScheduler | None = None,
+    start_loop: bool = True,
 ):
     """Construct a ``DiffusionEngine`` skeleton with a real busy loop.
 
     The engine is wired with a fake executor that asserts no concurrent
-    calls and a real ``RequestScheduler``.
+    calls and a real ``RequestScheduler``. ``start_loop=False`` leaves the
+    engine in its bootstrap state (no busy-loop thread).
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
     engine._closed = False
     engine.od_config = SimpleNamespace(streaming_output=False)
     engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
 
-    sched = RequestScheduler()
-    sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+    sched = scheduler
+    if sched is None:
+        sched = RequestScheduler()
+        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
     engine.scheduler = sched
     engine.step_execution = False
     engine.supports_request_batch = False
@@ -160,9 +181,14 @@ def _make_engine_with_loop(
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
+    engine._scheduling_paused = False
 
     engine.main_loop = loop
     engine.stop_event = threading.Event()
+    if not start_loop:
+        engine._loop_started = False
+        engine.worker_thread = None
+        return engine
     engine._loop_started = True
     engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True)
     engine.worker_thread.start()
@@ -170,6 +196,10 @@ def _make_engine_with_loop(
 
 
 def _stop_engine(engine: DiffusionEngine) -> None:
+    if engine.worker_thread is None:
+        return
+    for gate in engine.executor.gates.values():
+        gate.set()
     with engine._cv:
         engine.stop_event.set()
         engine._cv.notify_all()
@@ -565,3 +595,319 @@ async def test_rpc_and_request_results_do_not_swap():
     assert req_res.error == "result_for_req1"
     # Belt and braces: never overlapped on the executor.
     assert engine.executor.max_concurrent == 1
+
+
+# ───────────────────────── scheduler pause (keep) ─────────────────────────
+
+
+class _GatedAdmissionScheduler(RequestScheduler):
+    """Batching wait that only ends when the test opens admission."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.admission_open = threading.Event()
+
+    def get_admission_wait_decision(self, *, now: float, dp_concurrent: bool = False):
+        decision = super().get_admission_wait_decision(now=now, dp_concurrent=dp_concurrent)
+        return decision.__class__(should_wait=True, deadline=None, max_batch=self.max_num_running_reqs)
+
+    def should_end_admission_wait(self, decision, *, now: float, stable_since: float) -> bool:
+        return self.admission_open.is_set()
+
+
+def _pause(engine: DiffusionEngine, timeout: float | None = None):
+    return engine.async_collective_rpc("pause_scheduler", timeout=timeout, kwargs={"mode": "keep"})
+
+
+def _resume(engine: DiffusionEngine):
+    return engine.async_collective_rpc("resume_scheduler")
+
+
+def _methods(engine: DiffusionEngine) -> list[str]:
+    return [call["method"] for call in engine.executor.calls]
+
+
+def _executed(engine: DiffusionEngine) -> list[tuple[str, str]]:
+    return [
+        (call["args"][0].request_id, call["version"])
+        for call in engine.executor.calls
+        if call["method"] == "execute_model"
+    ]
+
+
+def _submit(engine: DiffusionEngine, tag: str) -> asyncio.Task:
+    return asyncio.create_task(_consume_final_output(engine.async_add_req_and_stream_response(_make_request(tag))))
+
+
+async def _wait_until(predicate, timeout: float = 3.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() > deadline:
+            raise AssertionError("condition not reached in time")
+        await asyncio.sleep(0.005)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_retains_waiting_request_and_serves_control_rpcs():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        assert await _pause(engine) is None
+        queued = _submit(engine, "queued")
+        result = await engine.async_collective_rpc("list_loras", unique_reply_rank=0)
+        assert result.error == "rpc_result_for_list_loras"
+        await asyncio.sleep(0.05)
+        assert engine.scheduler.num_waiting_requests() == 1
+        assert _executed(engine) == []
+        assert not queued.done()
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_queued"
+        assert _methods(engine) == ["synchronize_device", "list_loras", "execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_acks_after_running_batch_and_freezes_queue():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    gate = engine.executor.gates["A"] = threading.Event()
+    try:
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+
+        pause = asyncio.create_task(_pause(engine))
+        await asyncio.sleep(0.1)
+        assert not pause.done()
+
+        gate.set()
+        await asyncio.wait_for(pause, 3.0)
+        assert _methods(engine) == ["execute_model", "synchronize_device"]
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert engine.scheduler.num_waiting_requests() == 1
+        assert not b.done()
+
+        engine.executor.version = "v2"
+        await _resume(engine)
+        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
+        assert _executed(engine) == [("A", "v1"), ("B", "v2")]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_ack_does_not_release_later_rpcs_before_the_batch_is_emitted():
+    """A pause ACK must not open the queue to RPCs that touch device memory
+    (sleep) before the batch that just ran has been delivered.
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    barrier_gate = engine.executor.gates["synchronize_device"] = threading.Event()
+    methods_at_emit: list[list[str]] = []
+    emit_outputs = engine._emit_outputs
+
+    def _recording_emit(*args, **kwargs):
+        methods_at_emit.append(_methods(engine))
+        return emit_outputs(*args, **kwargs)
+
+    engine._emit_outputs = _recording_emit
+    try:
+        gate = engine.executor.gates["A"] = threading.Event()
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        pause = asyncio.create_task(_pause(engine))
+        await _wait_until(lambda: engine._scheduling_paused)
+        gate.set()
+        # The barrier holds the busy loop inside the post-execute RPC drain,
+        # so the sleep is queued before the pause ACK is published.
+        await _wait_until(lambda: "synchronize_device" in _methods(engine))
+        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
+        await _wait_until(lambda: not engine._rpc_queue.empty())
+        barrier_gate.set()
+
+        await asyncio.wait_for(pause, 3.0)
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
+    finally:
+        _stop_engine(engine)
+
+    assert _methods(engine) == ["execute_model", "synchronize_device", "handle_sleep_task"]
+    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
+
+
+@pytest.mark.asyncio
+async def test_timed_out_pause_still_ends_the_drain_before_the_batch_is_emitted():
+    """A pause whose caller already gave up is still a drain barrier: its
+    cancelled task must not let a later RPC run before the batch is emitted.
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    methods_at_emit: list[list[str]] = []
+    emit_outputs = engine._emit_outputs
+
+    def _recording_emit(*args, **kwargs):
+        methods_at_emit.append(_methods(engine))
+        return emit_outputs(*args, **kwargs)
+
+    engine._emit_outputs = _recording_emit
+    try:
+        gate = engine.executor.gates["A"] = threading.Event()
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+
+        # The caller times out while A is still executing: the pause future is
+        # cancelled, but its task stays queued ahead of the sleep.
+        with pytest.raises(TimeoutError):
+            await _pause(engine, timeout=0.1)
+        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
+        await _wait_until(lambda: engine._rpc_queue.qsize() == 2)
+        gate.set()
+
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
+    finally:
+        _stop_engine(engine)
+
+    assert _methods(engine) == ["execute_model", "handle_sleep_task"]
+    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
+
+
+@pytest.mark.asyncio
+async def test_pause_rejects_unsupported_modes_without_closing_the_gate():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        for kwargs in ({"mode": "abort"}, {"mode": "wait"}, {}):
+            with pytest.raises(ValueError):
+                await engine.async_collective_rpc("pause_scheduler", kwargs=kwargs)
+        engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
+        with pytest.raises(NotImplementedError):
+            await _pause(engine)
+        engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
+
+        assert (await asyncio.wait_for(_submit(engine, "still-runs"), 3.0)).error == "result_for_still-runs"
+        assert _methods(engine) == ["execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_during_batching_wait_holds_the_batch():
+    loop = asyncio.get_running_loop()
+    sched = _GatedAdmissionScheduler()
+    sched.initialize(SimpleNamespace(max_num_seqs=4, request_batch_max_wait_ms=1000.0))
+    engine = _make_engine_with_loop(loop, scheduler=sched)
+    try:
+        waiting = _submit(engine, "W")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+        await asyncio.sleep(0.05)
+
+        await asyncio.wait_for(_pause(engine), 3.0)
+        assert _methods(engine) == ["synchronize_device"]
+        assert engine.scheduler.num_waiting_requests() == 1
+
+        await _resume(engine)
+        await asyncio.sleep(0.05)
+        assert _executed(engine) == []
+        sched.admission_open.set()
+        assert (await asyncio.wait_for(waiting, 3.0)).error == "result_for_W"
+    finally:
+        sched.admission_open.set()
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_pause_timeout_leaves_gate_closed_until_explicit_resume():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    gate = engine.executor.gates["A"] = threading.Event()
+    try:
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+
+        with pytest.raises(TimeoutError):
+            await _pause(engine, timeout=0.1)
+        gate.set()
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        await asyncio.sleep(0.1)
+        assert _executed(engine) == [("A", "v1")]
+        assert not b.done()
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
+        assert "synchronize_device" not in _methods(engine)
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_repeated_pause_and_resume_are_idempotent():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        await _resume(engine)
+        await _pause(engine)
+        await _pause(engine)
+        queued = _submit(engine, "Q")
+        await asyncio.sleep(0.05)
+        assert _executed(engine) == []
+        await _resume(engine)
+        await _resume(engine)
+        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_Q"
+        assert _methods(engine) == ["synchronize_device", "synchronize_device", "execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_abort_of_waiting_request_while_paused_delivers_terminal():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        await _pause(engine)
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+        engine.abort("B")
+
+        result = await asyncio.wait_for(b, 3.0)
+        assert result.aborted is True
+        assert engine.scheduler.get_request_state("B") is None
+        assert _executed(engine) == []
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(_submit(engine, "C"), 3.0)).error == "result_for_C"
+        assert _executed(engine) == [("C", "v1")]
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
+def test_pause_and_resume_before_loop_starts_use_engine_local_path():
+    engine = _make_engine_with_loop(None, start_loop=False)
+
+    assert engine.collective_rpc("pause_scheduler", kwargs={"mode": "keep"}) is None
+    assert _methods(engine) == ["synchronize_device"]
+
+    assert engine.collective_rpc("resume_scheduler") is None
+    assert engine.add_req_and_wait_for_response(_make_request("X")).error == "result_for_X"
+    assert _executed(engine) == [("X", "v1")]
+
+
+@pytest.mark.asyncio
+async def test_requests_without_pause_never_trigger_the_barrier():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        first = await asyncio.wait_for(_submit(engine, "r1"), 3.0)
+        await engine.async_collective_rpc("ping", args=("p",), unique_reply_rank=0)
+        second = await asyncio.wait_for(_submit(engine, "r2"), 3.0)
+    finally:
+        _stop_engine(engine)
+
+    assert (first.error, second.error) == ("result_for_r1", "result_for_r2")
+    assert _methods(engine) == ["execute_model", "ping", "execute_model"]

tests/diffusion/test_diffusion_scheduler.py

@@ -2092,3 +2092,19 @@ class TestStepScheduler:
 
         with pytest.raises(ValueError):
             self.scheduler.add_request(request)
+
+
+class TestPendingFinishedRequestIds:
+    def test_reports_finished_ids_that_still_hold_state_without_clearing_them(self):
+        sched = RequestScheduler()
+        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+        sched.add_request(_make_request("a"))
+        sched.add_request(_make_request("b"))
+
+        sched.finish_requests("a", DiffusionRequestStatus.FINISHED_ABORTED)
+        assert sched.pending_finished_request_ids() == {"a"}
+
+        sched.pop_request_state("a")
+        assert sched.pending_finished_request_ids() == set()
+        # The next wave still ships the id to the worker for its own cleanup.
+        assert sched.schedule().finished_req_ids == {"a"}

tests/diffusion/test_diffusion_worker_synchronize_device.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for the device step of the diffusion pause barrier."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_worker(runner: object) -> DiffusionWorker:
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = runner
+    worker.rank = 0
+    return worker
+
+
+def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    order: list = []
+    manager = MagicMock()
+    manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True
+    platform.synchronize.side_effect = lambda: order.append("sync")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    worker.synchronize_device(timeout=2.5)
+
+    assert order == [("prefetch", 2.5), "sync"]
+
+
+def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    manager = MagicMock()
+    manager.wait_prefetch.return_value = False
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    with pytest.raises(TimeoutError, match="prefetch"):
+        worker.synchronize_device(timeout=0.01)
+
+    platform.synchronize.assert_not_called()
+
+
+def test_synchronize_device_without_transfer_manager_only_syncs(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=None))
+
+    worker.synchronize_device(timeout=None)
+
+    platform.synchronize.assert_called_once_with()

tests/distributed/omni_connectors/test_kv_async_prefetch.py

@@ -5,6 +5,7 @@ mock connector, dedup, per-call sender_info isolation, role classification, and
 miss/abort paths.  GPU H2D and D2D paths need CUDA (integration tests).
 """
 
+import threading
 from types import SimpleNamespace
 
 import pytest
@@ -354,3 +355,41 @@ def test_consume_then_apply_attaches_payload():
     receiver._record_stream_for_prefetched(data)
     receiver.apply_kv_cache_to_request(req, data)
     assert req.past_key_values is not None
+
+
+# --------------------------------------------------------------------------- #
+#  wait_prefetch: the pause barrier's view of outstanding background loads
+# --------------------------------------------------------------------------- #
+
+
+def test_wait_prefetch_is_true_when_nothing_outstanding():
+    _, receiver, _ = _make_sender_receiver()
+    assert receiver.wait_prefetch(timeout=0.01) is True
+
+
+def test_wait_prefetch_is_true_when_prefetch_disabled():
+    _, receiver, _ = _make_sender_receiver(async_prefetch=False)
+    assert receiver.wait_prefetch(timeout=0.01) is True
+
+
+def test_wait_prefetch_reports_a_running_prefetch_and_keeps_its_payload():
+    sender, receiver, connector = _make_sender_receiver()
+    _seed_payload(sender, "rid-w")
+    gate = threading.Event()
+    real_get = connector.get
+
+    def slow_get(*args, **kwargs):
+        gate.wait(timeout=5.0)
+        return real_get(*args, **kwargs)
+
+    connector.get = slow_get
+    receiver.start_prefetch({"request_id": "rid-w", "kv_sender_info": _SENDER_INFO})
+    try:
+        assert receiver.wait_prefetch(timeout=0.05) is False
+        gate.set()
+        assert receiver.wait_prefetch(timeout=5.0) is True
+        data, _ = receiver.consume_prefetched_kv(_req("rid-w"))
+        assert data is not None
+    finally:
+        gate.set()
+        receiver.shutdown_prefetch()

tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Lifecycle test for ``AsyncOmni.pause_generation(mode="keep")`` on a diffusion stage.
+
+Per round, request A is started, request B is queued behind it, and the pause
+is issued while A is still running. The pause must return only after A has
+been delivered, B must stay queued through ``sleep`` / ``wake_up``, and B may
+start only after ``resume_generation``.
+
+A second test follows the order the docs recommend for a weight update: sleep
+is issued as soon as the pause returns, without awaiting the request that was
+running, which must still be delivered intact.
+
+Both single-GPU executor backends are covered explicitly: ``uni`` (in-process
+worker) and ``mp`` (worker process with the asynchronous D2H output thread).
+
+Usage:
+    pytest tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py -v -s
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+from collections import defaultdict
+from contextlib import ExitStack
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.data import is_diffusion_request_started_output
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "tiny-random/Qwen-Image"
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
+
+# CI runs a couple of rounds; set PAUSE_KEEP_ROUNDS=10 for a soak run.
+ROUNDS = int(os.environ.get("PAUSE_KEEP_ROUNDS", "2"))
+# Enough denoising steps that the pause lands while A is still executing.
+NUM_INFERENCE_STEPS = 40
+
+
+def _sampling_params() -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        num_inference_steps=NUM_INFERENCE_STEPS,
+        guidance_scale=0.0,
+        height=256,
+        width=256,
+        seed=42,
+        emit_request_lifecycle=True,
+    )
+
+
+async def _generate(engine: AsyncOmni, request_id: str, events: dict[str, dict[str, float]]):
+    final = None
+    async for output in engine.generate(
+        prompt={"prompt_ids": list(range(50))},
+        request_id=request_id,
+        sampling_params_list=[_sampling_params()],
+        output_modalities=["image"],
+    ):
+        if is_diffusion_request_started_output(output):
+            events[request_id]["started"] = time.monotonic()
+            continue
+        final = output
+    events[request_id]["finished"] = time.monotonic()
+    return final
+
+
+async def _wait_for(predicate, timeout: float = 30.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        assert time.monotonic() < deadline, "condition not reached in time"
+        await asyncio.sleep(0.005)
+
+
+async def _pause_round(engine: AsyncOmni, events: dict[str, dict[str, float]], round_index: int) -> float:
+    """A runs, B queues, keep-pause lands while A runs; B starts only after resume.
+
+    Returns the pause latency.
+    """
+    a_id, b_id = f"a{round_index}", f"b{round_index}"
+    a = asyncio.create_task(_generate(engine, a_id, events))
+    await _wait_for(lambda: "started" in events[a_id])
+    b = asyncio.create_task(_generate(engine, b_id, events))
+    await asyncio.sleep(0.05)
+
+    assert not a.done(), "A finished before the pause was issued; raise NUM_INFERENCE_STEPS"
+    pause_start = time.monotonic()
+    await engine.pause_generation(mode="keep", clear_cache=False)
+    pause_latency = time.monotonic() - pause_start
+    assert await engine.is_paused()
+
+    a_output = await asyncio.wait_for(a, 60.0)
+    assert a_output.images, "A must be delivered, not aborted, by a keep pause"
+    assert "started" not in events[b_id], "B must not start while paused"
+
+    # Control RPCs and the sleep cycle stay available on the paused stage.
+    assert await engine.list_loras() == []
+    await engine.sleep(level=1)
+    assert await engine.is_sleeping()
+    await engine.wake_up()
+    assert not await engine.is_sleeping()
+    assert await engine.is_paused()
+    assert not b.done() and "started" not in events[b_id]
+
+    resume_at = time.monotonic()
+    await engine.resume_generation()
+    b_output = await asyncio.wait_for(b, 60.0)
+    assert b_output.images, "B must run after resume"
+    assert events[b_id]["started"] >= resume_at
+    return pause_latency
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["uni", "mp"])
+async def test_keep_pause_sleep_wake_resume_lifecycle(backend: str):
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            distributed_executor_backend=backend,
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+
+        pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(ROUNDS)]
+
+        assert not engine.request_states, "every request must be finalized"
+        print(
+            f"[{backend}] rounds={ROUNDS} steps={NUM_INFERENCE_STEPS} "
+            f"pause_latency_s min/median/max="
+            f"{min(pause_latencies):.3f}/{sorted(pause_latencies)[ROUNDS // 2]:.3f}/{max(pause_latencies):.3f}"
+        )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["uni", "mp"])
+async def test_sleep_right_after_keep_pause_keeps_the_running_batch(backend: str):
+    """The documented weight-update order: sleep is issued as soon as the pause
+    returns, with the request that was running still un-awaited.
+    """
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            distributed_executor_backend=backend,
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+
+        a = asyncio.create_task(_generate(engine, "a", events))
+        await _wait_for(lambda: "started" in events["a"])
+        await asyncio.sleep(0.05)
+        assert not a.done(), "A finished before the pause was issued; raise NUM_INFERENCE_STEPS"
+
+        await engine.pause_generation(mode="keep", clear_cache=False)
+        await engine.sleep(level=1)
+
+        a_output = await asyncio.wait_for(a, 60.0)
+        assert a_output.images, "A must survive a sleep issued right after the pause ACK"
+
+        await engine.wake_up()
+        await engine.resume_generation()
+        b_output = await asyncio.wait_for(_generate(engine, "b", events), 120.0)
+        assert b_output.images, "the stage must generate again after wake and resume"

tests/e2e/features/helpers/custom_pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Test doubles for custom-pipeline E2E and RLHF integration.
 
@@ -290,6 +290,36 @@ class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
         """Return a stable identifier for assertions in unit tests."""
         return "vllm-omni-colocate-worker-extension-for-test"
 
+    def start_counting_synchronize_device(self) -> bool:
+        """Count pause-barrier runs on this rank (test-only)."""
+        self._synchronize_device_calls = 0
+        inner = type(self).synchronize_device
+
+        def _counted(timeout=None):
+            self._synchronize_device_calls += 1
+            return inner(self, timeout=timeout)
+
+        self.synchronize_device = _counted
+        return True
+
+    def synchronize_device_seen(self) -> bool:
+        """True from every rank only when each rank ran the barrier at least once."""
+        return getattr(self, "_synchronize_device_calls", 0) > 0
+
+    def set_synchronize_device_failure(self, rank: int, fail: bool) -> bool:
+        """Make (or stop making) the pause barrier fail on ``rank`` (test-only)."""
+        if self.rank != rank:
+            return True
+        if fail:
+
+            def _fail(timeout=None):
+                raise RuntimeError(f"injected barrier failure on rank {rank}")
+
+            self.synchronize_device = _fail
+        else:
+            self.__dict__.pop("synchronize_device", None)
+        return True
+
 
 # ---------------------------------------------------------------------------
 # Pipeline layer

tests/engine/test_orchestrator_error_handling.py

@@ -13,6 +13,7 @@ import asyncio
 import queue
 import time
 from types import SimpleNamespace
+from unittest.mock import AsyncMock
 
 import janus
 import pytest
@@ -22,6 +23,7 @@ from vllm.v1.serial_utils import MsgpackEncoder
 
 from vllm_omni.engine.messages import (
     AddCompanionRequestMessage,
+    CollectiveRPCRequestMessage,
     EngineQueueMessage,
     ErrorMessage,
     ShutdownRequestMessage,
@@ -1143,3 +1145,64 @@ async def test_diffusion_client_error_output_propagates_non_400_status(
     finally:
         orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
         orchestrator_fixture.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal() -> None:
+    """A control RPC that fails on one replica must come back as that
+    replica's error result instead of raising out of _request_handler.
+    """
+
+    class _FailingClient(FakeStageClient):
+        async def collective_rpc_async(self, *args, **kwargs):
+            raise TimeoutError("outputs did not drain")
+
+    ok = FakeStageClient(stage_type="diffusion")
+    ok.collective_rpc_async = AsyncMock(return_value=None)
+    orchestrator, queues = _build_bare_orchestrator(
+        _build_stage_pools([[ok], [_FailingClient(stage_type="diffusion")]])
+    )
+    try:
+        await orchestrator._handle_collective_rpc(
+            CollectiveRPCRequestMessage(
+                rpc_id="rpc-1",
+                method="pause_scheduler",
+                args=(),
+                kwargs={"mode": "keep"},
+                stage_ids=[0, 1],
+            )
+        )
+
+        result = queues[2].async_q.get_nowait()
+        assert result.rpc_id == "rpc-1"
+        assert result.stage_ids == [0, 1]
+        assert result.results[0] is None
+        assert result.results[1]["supported"] is False
+        assert "outputs did not drain" in result.results[1]["error"]
+    finally:
+        for q in queues:
+            q.close()
+
+
+@pytest.mark.asyncio
+async def test_rpc_failure_capture_is_limited_to_pause_and_resume() -> None:
+    """Only pause/resume failures are reported as a result; every other RPC
+    keeps propagating the way it does today.
+    """
+    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="diffusion")]]))
+    orchestrator.stage_pools[0].collective_rpc = AsyncMock(side_effect=TimeoutError("worker died"))
+
+    def _msg(rpc_id: str, method: str) -> CollectiveRPCRequestMessage:
+        return CollectiveRPCRequestMessage(rpc_id=rpc_id, method=method, args=(), kwargs={}, stage_ids=[0])
+
+    try:
+        await orchestrator._handle_collective_rpc(_msg("rpc-resume", "resume_scheduler"))
+        result = queues[2].async_q.get_nowait()
+        assert result.results[0]["supported"] is False
+        assert "worker died" in result.results[0]["error"]
+
+        with pytest.raises(TimeoutError):
+            await orchestrator._handle_collective_rpc(_msg("rpc-sleep", "sleep"))
+    finally:
+        for q in queues:
+            q.close()

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -25,6 +25,7 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
     omni._paused = False
     omni._admitting = 0
     omni._hold_admission_until_resume = False
+    omni._paused_stage_ids = set()
     omni._sleeping_tags = set()
     omni._stage_sleeping_tags = {}
     omni._level2_sleeping = False
@@ -460,3 +461,225 @@ def test_partial_wake_does_not_skip_remaining_sleeping_stage():
         assert not omni._stage_sleeping_tags
 
     asyncio.run(run())
+
+
+# ───────────────────── diffusion pause_generation(mode="keep") ─────────────────────
+
+
+def _rpc_methods(omni: AsyncOmni) -> list[tuple[str, list[int] | None]]:
+    return [(call.kwargs["method"], call.kwargs.get("stage_ids")) for call in omni.collective_rpc.await_args_list]
+
+
+@pytest.mark.cpu
+def test_diffusion_keep_routes_to_backend():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+
+        await omni.pause_generation(mode="keep", clear_cache=False)
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="pause_scheduler",
+            args=(),
+            kwargs={"mode": "keep"},
+            stage_ids=[0],
+        )
+        assert await omni.is_paused() is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "pause_kwargs",
+    [
+        {},
+        {"mode": "abort"},
+        {"mode": "wait"},
+        {"mode": "keep", "wait_for_inflight_requests": True},
+    ],
+)
+def test_diffusion_non_keep_modes_pause_frontend_only(pause_kwargs):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+
+        await omni.pause_generation(**pause_kwargs)
+        assert await omni.is_paused() is True
+        omni.collective_rpc.assert_not_awaited()
+
+        await omni.resume_generation()
+        omni.collective_rpc.assert_not_awaited()
+        assert await omni.is_paused() is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_mixed_engine_keep_sends_engine_core_and_diffusion_pauses():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+
+        await omni.pause_generation(mode="keep", clear_cache=True)
+
+        assert omni.collective_rpc.await_args_list[0].kwargs == {
+            "method": "pause_scheduler",
+            "args": (),
+            "kwargs": {"mode": "keep", "clear_cache": True},
+            "stage_ids": [0],
+        }
+        assert omni.collective_rpc.await_args_list[1].kwargs == {
+            "method": "pause_scheduler",
+            "args": (),
+            "kwargs": {"mode": "keep"},
+            "stage_ids": [1],
+        }
+        # The frontend cache reset still runs after the backend pauses.
+        omni.reset_prefix_cache.assert_awaited_once()
+
+        await omni.resume_generation()
+        assert _rpc_methods(omni)[2:] == [("resume_scheduler", [0]), ("resume_scheduler", [1])]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_keep_targets_only_requested_diffusion_stages():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion", "diffusion"])
+
+        await omni.pause_generation(mode="keep", clear_cache=False, stage_ids=[1])
+        assert _rpc_methods(omni) == [("pause_scheduler", [1])]
+
+        await omni.resume_generation(stage_ids=[0])
+        assert _rpc_methods(omni) == [("pause_scheduler", [1])]
+        # Stage 1's scheduler is still closed: admission must stay closed too,
+        # otherwise new requests queue on a stage that will not schedule them.
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation(stage_ids=[1])
+        assert _rpc_methods(omni) == [("pause_scheduler", [1]), ("resume_scheduler", [1])]
+        assert await omni.is_paused() is False
+
+        await omni.resume_generation()
+        assert len(_rpc_methods(omni)) == 2
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_keep_then_sleep_wake_resume_keeps_admission_paused_until_resume():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_encoder_cache = AsyncMock()
+        omni.collective_rpc = AsyncMock(
+            side_effect=lambda **kw: [OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)]
+            if kw["method"] in {"handle_sleep_task", "handle_wake_task"}
+            else [None]
+        )
+
+        await omni.pause_generation(mode="keep", clear_cache=False)
+        await omni.sleep(level=1)
+        await omni.wake_up()
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation()
+        assert await omni.is_paused() is False
+        assert [m for m, _ in _rpc_methods(omni)] == [
+            "pause_scheduler",
+            "handle_sleep_task",
+            "handle_wake_task",
+            "resume_scheduler",
+        ]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "failure",
+    [
+        AsyncMock(side_effect=NotImplementedError("request-level execution only")),
+        AsyncMock(return_value=[{"supported": False, "error": "TimeoutError: outputs did not drain"}]),
+    ],
+)
+def test_keep_failure_surfaces_and_leaves_explicit_resume_possible(failure):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.collective_rpc = failure
+
+        with pytest.raises((NotImplementedError, RuntimeError)):
+            await omni.pause_generation(mode="keep", clear_cache=False)
+
+        assert await omni.is_paused() is True
+        assert omni._paused_stage_ids == {0}
+
+        omni.collective_rpc = AsyncMock(return_value=[None])
+        await omni.resume_generation()
+        omni.collective_rpc.assert_awaited_once_with(
+            method="resume_scheduler",
+            args=(),
+            kwargs=None,
+            stage_ids=[0],
+        )
+        assert await omni.is_paused() is False
+        assert omni._paused_stage_ids == set()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_resume_failure_keeps_admission_paused_and_target_recorded():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        await omni.pause_generation(mode="keep", clear_cache=False)
+
+        omni.collective_rpc = AsyncMock(return_value=[{"error": "replica gone"}])
+        with pytest.raises(RuntimeError, match="resume_scheduler failed"):
+            await omni.resume_generation()
+
+        assert await omni.is_paused() is True
+        assert omni._paused_stage_ids == {0}
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("first, second", [([1], [0]), ([0], [1])])
+def test_mixed_engine_partial_resume_keeps_admission_paused(first, second):
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        await omni.pause_generation(mode="keep", clear_cache=False)
+
+        await omni.resume_generation(stage_ids=first)
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation(stage_ids=second)
+        assert await omni.is_paused() is False
+        assert _rpc_methods(omni)[2:] == [("resume_scheduler", first), ("resume_scheduler", second)]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_two_ar_stages_partial_resume_keeps_admission_paused():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "llm"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+        await omni.pause_generation(mode="abort")
+
+        await omni.resume_generation(stage_ids=[0])
+        assert await omni.is_paused() is True
+
+        await omni.resume_generation(stage_ids=[1])
+        assert await omni.is_paused() is False
+
+    asyncio.run(run())

vllm_omni/diffusion/diffusion_engine.py

@@ -228,9 +228,10 @@ class DiffusionEngine:
     #: remains a pure explicit user override (never mutated by engines).
     default_diffusion_model_runner_cls: str | None = None
 
-    # Class-level default so tests using object.__new__ (without __init__)
-    # don't hit AttributeError when _busy_loop accesses self.dp_concurrent.
+    # Class-level defaults so tests using object.__new__ (without __init__)
+    # don't hit AttributeError when _busy_loop accesses them.
     dp_concurrent: bool = False
+    _scheduling_paused: bool = False
 
     def __init__(
         self,
@@ -387,6 +388,8 @@ class DiffusionEngine:
         self._shutdown_complete = False
         self.abort_queue: queue.Queue[str] = queue.Queue()
         self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
+        # pause_scheduler(mode="keep"): no new batch is scheduled while set.
+        self._scheduling_paused = False
         # Copied onto the existing output metrics payload so queue monitoring
         # reuses the normal diffusion result path without additional IPC.
         self._scheduler_num_waiting_reqs = 0
@@ -584,10 +587,16 @@ class DiffusionEngine:
         while not self.stop_event.is_set():
             self._process_aborts_queue()
             self._process_rpc_queue()
+            if self._scheduling_paused:
+                # No wave runs while paused, so requests finished by an abort
+                # would otherwise wait for the resume to surface.
+                with self._cv:
+                    pending_finished = self.scheduler.pending_finished_request_ids()
+                self._emit_finished_outputs(pending_finished, None)
 
             with self._cv:
                 while (
-                    not self.scheduler.has_requests()
+                    (self._scheduling_paused or not self.scheduler.has_requests())
                     and self._rpc_queue.empty()
                     and self.abort_queue.empty()
                     and not self.stop_event.is_set()
@@ -597,11 +606,14 @@ class DiffusionEngine:
                 if self.stop_event.is_set():
                     break
 
-                if not self.scheduler.has_requests():
+                if self._scheduling_paused or not self.scheduler.has_requests():
                     # Only RPC / abort work pending; loop back to drain it.
                     continue
 
                 self._wait_for_admission_if_needed_locked()
+                if self._scheduling_paused:
+                    # The gate closed while the batching wait released the lock.
+                    continue
 
                 try:
                     sched_output = self.scheduler.schedule()
@@ -693,7 +705,7 @@ class DiffusionEngine:
         last_waiting = -1
         stable_since = start
 
-        while not self.stop_event.is_set():
+        while not self.stop_event.is_set() and not self._scheduling_paused:
             waiting = self.scheduler.num_waiting_requests()
             now = time.monotonic()
 
@@ -733,19 +745,34 @@ class DiffusionEngine:
             except queue.Empty:
                 return
 
-            fut = task.future
-            if fut.cancelled() or fut.done():
-                continue
+            self._run_rpc_task(task)
 
-            remaining: float | None = None
-            if task.deadline is not None:
-                remaining = task.deadline - time.monotonic()
-                if remaining <= 0:
-                    if not fut.done():
-                        fut.set_exception(TimeoutError(f"RPC call to {task.method} timed out before execution."))
-                    continue
+            if task.method == "pause_scheduler":
+                # A dequeued pause ends this drain whatever became of it, so
+                # the batch that just ran is delivered before anything the
+                # caller queued behind the pause. The rest of the queue is
+                # picked up on the next pass of the busy loop.
+                return
 
-            try:
+    def _run_rpc_task(self, task: _RpcTask) -> None:
+        fut = task.future
+        if fut.cancelled() or fut.done():
+            return
+
+        remaining: float | None = None
+        if task.deadline is not None:
+            remaining = task.deadline - time.monotonic()
+            if remaining <= 0:
+                if not fut.done():
+                    fut.set_exception(TimeoutError(f"RPC call to {task.method} timed out before execution."))
+                return
+
+        try:
+            if task.method == "pause_scheduler":
+                result = self._run_pause_barrier(remaining)
+            elif task.method == "resume_scheduler":
+                result = self._open_scheduling_gate()
+            else:
                 result = self.executor.collective_rpc(
                     method=task.method,
                     timeout=remaining,
@@ -753,16 +780,16 @@ class DiffusionEngine:
                     kwargs=task.kwargs,
                     unique_reply_rank=task.unique_reply_rank,
                 )
-            except BaseException as exc:  # noqa: BLE001 - propagate to caller
-                # The future may have been cancelled (e.g. by a sync timeout
-                # or asyncio cancellation) while the executor call was
-                # running. Setting state on a cancelled/done future raises
-                # InvalidStateError, which would kill the busy loop.
-                if not fut.done():
-                    fut.set_exception(exc)
-            else:
-                if not fut.done():
-                    fut.set_result(result)
+        except BaseException as exc:  # noqa: BLE001 - propagate to caller
+            # The future may have been cancelled (e.g. by a sync timeout
+            # or asyncio cancellation) while the executor call was
+            # running. Setting state on a cancelled/done future raises
+            # InvalidStateError, which would kill the busy loop.
+            if not fut.done():
+                fut.set_exception(exc)
+        else:
+            if not fut.done():
+                fut.set_result(result)
 
     def _fail_pending_rpcs(self, exc: BaseException) -> None:
         while True:
@@ -773,6 +800,30 @@ class DiffusionEngine:
             if not task.future.done():
                 task.future.set_exception(exc)
 
+    def _check_pause_request(self, kwargs: dict | None) -> None:
+        mode = (kwargs or {}).get("mode")
+        if mode != "keep":
+            raise ValueError(f"DiffusionEngine pause supports mode='keep' only, got {mode!r}.")
+        if self.execution_mode != DiffusionExecutionMode.REQUEST_BATCH:
+            raise NotImplementedError("DiffusionEngine pause supports request-level execution only.")
+
+    def _run_pause_barrier(self, timeout: float | None) -> None:
+        """Return once every worker has finished the device work of the batch that ran."""
+        self.executor.collective_rpc("synchronize_device", timeout=timeout, kwargs={"timeout": timeout})
+
+    def _open_scheduling_gate(self) -> None:
+        with self._cv:
+            self._scheduling_paused = False
+            self._cv.notify_all()
+
+    def _run_engine_control(self, method: str, timeout: float | None, kwargs: dict | None) -> None:
+        """Bootstrap-path dispatch for engine-local control methods."""
+        if method == "pause_scheduler":
+            self._check_pause_request(kwargs)
+            self._scheduling_paused = True
+            return self._run_pause_barrier(timeout)
+        return self._open_scheduling_gate()
+
     def _remove_diffusion_kv_requests(self, request_ids: Iterable[str]) -> None:
         """Clear terminal Worker rows while Scheduler owns the allocations."""
 
@@ -1309,6 +1360,9 @@ class DiffusionEngine:
             unique_reply_rank=unique_reply_rank,
         )
         with self._cv:
+            if method == "pause_scheduler":
+                self._check_pause_request(kwargs)
+                self._scheduling_paused = True
             self._rpc_queue.put(task)
             self._cv.notify_all()
         return task
@@ -1351,6 +1405,8 @@ class DiffusionEngine:
                 # between the outer check and acquiring the lock, in which
                 # case we should use the queued path for proper ordering.
                 if not self._loop_started:
+                    if method in ("pause_scheduler", "resume_scheduler"):
+                        return self._run_engine_control(method, timeout, kwargs)
                     return self.executor.collective_rpc(
                         method=method,
                         timeout=timeout,

vllm_omni/diffusion/sched/base_scheduler.py

@@ -343,6 +343,10 @@ class BaseScheduler(ABC):
     def num_running_requests(self) -> int:
         return len(self._running)
 
+    def pending_finished_request_ids(self) -> set[str]:
+        """Finished requests whose state the engine has not consumed yet."""
+        return {request_id for request_id in self._finished_req_ids if request_id in self._request_states}
+
     def get_admission_wait_decision(
         self,
         *,

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -877,6 +877,19 @@ class DiffusionWorker:
         logger.info(f"[Worker {self.rank}] Wake-up complete.")
         return True
 
+    def synchronize_device(self, timeout: float | None = None) -> None:
+        """Wait until this rank has no device work left.
+
+        A KV prefetch still receiving on its background thread has queued no
+        device work yet, so it is joined first. ``timeout`` bounds that join
+        (and the multi-process worker's output drain before this call); the
+        device wait itself is unbounded.
+        """
+        manager = getattr(self.model_runner, "kv_transfer_manager", None)
+        if manager is not None and not manager.wait_prefetch(timeout=timeout):
+            raise TimeoutError("Diffusion KV prefetch did not finish before pause")
+        current_omni_platform.synchronize()
+
     def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK | None:
         from vllm_omni.platforms import current_omni_platform
 
@@ -1214,11 +1227,13 @@ class WorkerProc:
                     self._async_output_pending = max(0, self._async_output_pending - 1)
                     self._async_output_done.notify_all()
 
-    def drain_async_outputs(self, timeout: float = _ASYNC_OUTPUT_DRAIN_TIMEOUT_S) -> bool:
+    def drain_async_outputs(self, timeout: float | None = None) -> bool:
         """Block until background D2H/SHM packing has no work left.
 
         Returns False if outputs are still in flight when ``timeout`` expires.
         """
+        if timeout is None:
+            timeout = _ASYNC_OUTPUT_DRAIN_TIMEOUT_S
         with self._async_output_done:
             if self._async_output_pending == 0:
                 return True
@@ -1334,6 +1349,8 @@ class WorkerProc:
         }
 
         try:
+            if method == "synchronize_device" and not self.drain_async_outputs(timeout=kwargs.get("timeout")):
+                raise TimeoutError("Diffusion async outputs did not drain before pause")
             if method in _MEMORY_RELEASING_METHODS:
                 self.drain_async_outputs()
             # Use execute_method from WorkerWrapperBase for consistent method resolution

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -8,6 +8,7 @@ import struct
 import time
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import wait as wait_futures
 from dataclasses import asdict, dataclass
 from typing import Any
 
@@ -1346,6 +1347,16 @@ class OmniKVTransferManager:
             logger.exception("KV load failed for %s; falling back to sync receive", request_id)
             return None, 0
 
+    def wait_prefetch(self, timeout: float | None = None) -> bool:
+        """Block until no background prefetch is running; payloads stay consumable.
+
+        Returns False if a prefetch is still in flight when ``timeout`` expires.
+        """
+        futures = list(self._prefetch_futures.values())
+        if not futures:
+            return True
+        return not wait_futures(futures, timeout=timeout).not_done
+
     def _discard_future(self, request_id: str) -> None:
         """Cancel an unstarted prefetch or attach a callback to drop a running one."""
         fut = self._prefetch_futures.pop(request_id, None)

vllm_omni/engine/orchestrator.py

@@ -696,13 +696,23 @@ class OrchestratorBase:
         stage_ids: list[int] = []
         for pool in target_pools:
             for replica_id in pool.live_replica_ids():
-                stage_result = await pool.collective_rpc(
-                    replica_id=replica_id,
-                    method=method,
-                    timeout=timeout,
-                    args=args,
-                    kwargs=kwargs,
-                )
+                try:
+                    stage_result = await pool.collective_rpc(
+                        replica_id=replica_id,
+                        method=method,
+                        timeout=timeout,
+                        args=args,
+                        kwargs=kwargs,
+                    )
+                except Exception as exc:
+                    if method not in ("pause_scheduler", "resume_scheduler"):
+                        raise
+                    # A pause or resume that fails on one replica leaves the
+                    # engine usable, so report it as this replica's result
+                    # rather than out of the request handler; the caller
+                    # reaches both through _engine_core_rpc, which raises on
+                    # the error result and can then retry or resume.
+                    stage_result = {"supported": False, "error": f"{type(exc).__name__}: {exc}"}
                 stage_ids.append(pool.stage_id)
                 results.append(stage_result)
 

vllm_omni/entrypoints/async_omni.py

@@ -95,6 +95,10 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         # sleep uses _paused as a temporary admission gate and clears it
         # on wake so sleep → wake → generate keeps working.
         self._hold_admission_until_resume: bool = False
+        # Stages whose scheduler pause_generation closed (AR in any mode,
+        # diffusion with mode="keep"); admission stays closed until
+        # resume_generation has reopened all of them.
+        self._paused_stage_ids: set[int] = set()
         self._sleeping_tags: set[str] = set()
         self._stage_sleeping_tags: dict[int, set[str]] = {}
         self._level2_sleeping: bool = False
@@ -542,10 +546,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ) -> list[Any]:
-        """Call an AR EngineCore helper via collective_rpc (orchestrator loop).
+        """Call an engine control helper via collective_rpc (orchestrator loop).
 
         StagePool resolves ``{method}_async`` on the AR client when present
-        (vLLM AsyncMPClient convention). Raises if any replica reports failure.
+        (vLLM AsyncMPClient convention); diffusion stages answer the same
+        method names inside DiffusionEngine. Raises if any replica reports
+        failure.
         """
         results = await self.collective_rpc(
             method=method,
@@ -695,8 +701,13 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         1. Stop frontend admission (``_paused``).
         2. For AR/LLM stages, call EngineCore.pause_scheduler via the
            Orchestrator loop (abort/wait/keep + optional cache clear).
-        3. Diffusion stages have no EngineCore scheduler — only frontend
-           admission is paused for them.
+        3. For diffusion stages, ``mode="keep"`` pauses the DiffusionEngine
+           scheduler and returns once the batch that was running has finished
+           on every worker; that batch is delivered before any control RPC
+           issued after this call runs, so the documented pause -> sleep order
+           is safe. Queued requests stay queued until
+           :meth:`resume_generation`. Other modes pause frontend admission
+           only.
 
         Note: ``sleep()`` already pauses the AR scheduler internally (same as
         vLLM EngineCore.sleep). Call this API when you need pause *without*
@@ -711,7 +722,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             self._paused = True
             self._hold_admission_until_resume = True
 
-        ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        if mode != "keep":
+            diffusion_stage_ids = []
+        # Recorded before the RPCs so a failed or cancelled pause can still
+        # be undone with an explicit resume_generation.
+        self._paused_stage_ids.update(ar_stage_ids, diffusion_stage_ids)
         if ar_stage_ids:
             logger.info(
                 "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)",
@@ -726,6 +742,17 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                 stage_ids=ar_stage_ids,
                 kwargs={"mode": mode, "clear_cache": clear_cache},
             )
+        if diffusion_stage_ids:
+            logger.info(
+                "[%s] Pausing diffusion stage(s) %s via DiffusionEngine pause_scheduler(mode=keep)",
+                self._name,
+                diffusion_stage_ids,
+            )
+            await self._engine_core_rpc(
+                "pause_scheduler",
+                stage_ids=diffusion_stage_ids,
+                kwargs={"mode": "keep"},
+            )
 
         # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler
         # already clears AR-side caches when clear_cache=True.
@@ -739,10 +766,26 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
 
     async def resume_generation(self, stage_ids: list[int] | None = None) -> None:
         """Resume generation after :meth:`pause_generation`."""
-        ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
         if ar_stage_ids:
             logger.info("[%s] Resuming AR stage(s) %s via EngineCore", self._name, ar_stage_ids)
             await self._engine_core_rpc("resume_scheduler", stage_ids=ar_stage_ids)
+            self._paused_stage_ids.difference_update(ar_stage_ids)
+        diffusion_stage_ids = [sid for sid in diffusion_stage_ids if sid in self._paused_stage_ids]
+        if diffusion_stage_ids:
+            logger.info("[%s] Resuming diffusion stage(s) %s via DiffusionEngine", self._name, diffusion_stage_ids)
+            await self._engine_core_rpc("resume_scheduler", stage_ids=diffusion_stage_ids)
+            self._paused_stage_ids.difference_update(diffusion_stage_ids)
+
+        if self._paused_stage_ids:
+            # Reopening admission now would let new requests queue on a stage
+            # whose scheduler is still closed.
+            logger.info(
+                "[%s] Admission stays paused: stage(s) %s are still paused",
+                self._name,
+                sorted(self._paused_stage_ids),
+            )
+            return
 
         async with self._pause_cond:
             self._paused = False
@@ -816,8 +859,9 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         AR/LLM stages use EngineCore.sleep (pause scheduler, wait idle, then
         offload/discard memory) — matching vLLM AsyncLLM.sleep.
 
-        Diffusion stages keep the existing worker-level handle_sleep_task RPC
-        because StageDiffusionProc does not expose EngineCore.pause_scheduler.
+        Diffusion stages keep the worker-level handle_sleep_task RPC, which
+        does not stop the DiffusionEngine scheduler; quiesce a busy diffusion
+        stage first with ``pause_generation(mode="keep")`` (or abort it).
 
         Frontend admission is blocked at the start of this call (``_paused``)
         so pipelined :meth:`generate` cannot race into stages while sleep is