[Feature]: Support pause / resume and sleep / wake for AR stages in AsyncOmni (#6084)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/engine/test_async_omni_engine_abort_ack.py

@@ -0,0 +1,156 @@
+"""Unit tests for acknowledged abort on AsyncOmniEngine."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+from pytest_mock import MockerFixture
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.messages import AbortRequestMessage, AbortResultMessage
+from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine:
+    engine = object.__new__(AsyncOmniEngine)
+    engine.request_queue = SimpleNamespace(sync_q=request_q)
+    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
+    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
+    return engine
+
+
+def test_abort_result_message_correlation_key():
+    msg = AbortResultMessage(rpc_id="abc", success=True)
+    assert msg.rpc_correlation_key == ("abort", "abc")
+
+
+def test_sync_abort_is_fire_and_forget():
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+
+    engine.abort(["req-1"])
+
+    msg = request_q.get_nowait()
+    assert isinstance(msg, AbortRequestMessage)
+    assert msg.request_ids == ["req-1"]
+    assert msg.rpc_id is None
+    assert rpc_q.empty()
+    engine._correlated_rpc_client.close()
+
+
+def test_abort_async_waits_for_ack(mocker: MockerFixture):
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="abort-rpc-1"),
+    )
+
+    async def _run() -> None:
+        task = asyncio.create_task(engine.abort_async(["req-a", "req-b"], timeout=1))
+        msg = await asyncio.to_thread(request_q.get, True, 1)
+        assert isinstance(msg, AbortRequestMessage)
+        assert msg.request_ids == ["req-a", "req-b"]
+        assert msg.rpc_id == "abort-rpc-1"
+        assert not task.done()
+
+        rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
+        await task
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()
+
+
+def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture):
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="abort-rpc-err"),
+    )
+
+    async def _run() -> None:
+        task = asyncio.create_task(engine.abort_async(["req-x"], timeout=1))
+        await asyncio.to_thread(request_q.get, True, 1)
+        rpc_q.put(
+            AbortResultMessage(
+                rpc_id="abort-rpc-err",
+                success=False,
+                error="stage abort failed",
+            )
+        )
+        with pytest.raises(RuntimeError, match="stage abort failed"):
+            await task
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()
+
+
+def test_abort_async_times_out_without_result(mocker: MockerFixture):
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="abort-rpc-timeout"),
+    )
+
+    async def _run() -> None:
+        with pytest.raises(TimeoutError, match="abort timed out"):
+            await engine.abort_async(["req-timeout"], timeout=0.05)
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()
+
+
+def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture):
+    class SignallingQueue(queue.Queue):
+        def __init__(self) -> None:
+            super().__init__(maxsize=1)
+            self.put_attempted = threading.Event()
+
+        def put(self, item, block=True, timeout=None):
+            self.put_attempted.set()
+            return super().put(item, block=block, timeout=timeout)
+
+    request_q = SignallingQueue()
+    request_q.put("queue-is-full")
+    request_q.put_attempted.clear()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="blocked-abort"),
+    )
+
+    async def _run() -> None:
+        task = asyncio.create_task(engine.abort_async(["req-blocked"], timeout=1))
+        assert await asyncio.to_thread(request_q.put_attempted.wait, 1)
+        assert not task.done()
+
+        assert request_q.get(timeout=1) == "queue-is-full"
+        msg = await asyncio.to_thread(request_q.get, True, 1)
+        assert isinstance(msg, AbortRequestMessage)
+        assert msg.rpc_id == "blocked-abort"
+        rpc_q.put(AbortResultMessage(rpc_id="blocked-abort", success=True))
+        await task
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()

tests/engine/test_orchestrator.py

@@ -19,6 +19,7 @@ from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -805,10 +806,74 @@ async def test_run_abort(orchestrator_factory) -> None:
         assert stages[0].abort_calls == [["req-abort"]]
         assert stages[1].abort_calls == []
         assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
+        # Fire-and-forget abort must not emit an RPC result.
+        assert orchestrator_fixture.queues[2].sync_q.empty()
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
 
 
+@pytest.mark.asyncio
+async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> None:
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=False),
+        FakeStageClient(stage_type="llm", final_output=True),
+    ]
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[2], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory(stages, output_processors=processors)
+    request = SimpleNamespace(request_id="req-ack", prompt_token_ids=[1, 2, 3])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-ack",
+            prompt=request,
+            original_prompt={"prompt": "cancel with ack"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+
+        orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["req-ack"], rpc_id="abort-1"))
+        result = await _get_rpc_message(orchestrator_fixture)
+        assert isinstance(result, AbortResultMessage)
+        assert result.rpc_id == "abort-1"
+        assert result.success is True
+        assert result.error is None
+        assert result.rpc_correlation_key == ("abort", "abort-1")
+        assert stages[0].abort_calls == [["req-ack"]]
+        assert "req-ack" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_handle_abort_emits_error_result_on_failure() -> None:
+    rpc_queue: asyncio.Queue = asyncio.Queue()
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=rpc_queue,
+        stage_pools=[],
+    )
+
+    async def boom(_request_ids, *, abort=False):
+        del _request_ids, abort
+        raise RuntimeError("cleanup exploded")
+
+    orchestrator._cleanup_request_ids = boom  # type: ignore[method-assign]
+
+    await orchestrator._handle_abort(AbortRequestMessage(request_ids=["req-fail"], rpc_id="abort-err"))
+
+    result = rpc_queue.get_nowait()
+    assert isinstance(result, AbortResultMessage)
+    assert result.rpc_id == "abort-err"
+    assert result.success is False
+    assert "cleanup exploded" in (result.error or "")
+
+
 def _duplex_open_message(
     session_id: str,
     *,

tests/engine/test_rpc_result_router.py

@@ -4,6 +4,7 @@ from typing import Literal
 import pytest
 
 from vllm_omni.engine.messages import (
+    AbortResultMessage,
     CollectiveRPCResultMessage,
     EngineQueueMessage,
     ErrorMessage,
@@ -93,6 +94,19 @@ def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors():
     router.close()
 
 
+def test_rpc_result_router_routes_abort_results_by_correlation_id():
+    source: queue.Queue = queue.Queue()
+    router = RpcResultRouter(source)
+    waiter = router.register(("abort", "abort-42"))
+
+    source.put(AbortResultMessage(rpc_id="abort-42", success=True))
+
+    result = waiter.get(timeout=1)
+    assert isinstance(result, AbortResultMessage)
+    assert result.rpc_id == "abort-42"
+    router.close()
+
+
 def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer():
     source: queue.Queue = queue.Queue()
     router = RpcResultRouter(source)

tests/engine/test_stage_pool_collective_rpc.py

@@ -0,0 +1,71 @@
+"""Unit tests for StagePool.collective_rpc EngineCore control dispatch."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from vllm_omni.engine.stage_pool import StagePool
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_pool(*, stage_type: str = "llm", **client_methods: AsyncMock) -> tuple[StagePool, SimpleNamespace]:
+    client = SimpleNamespace(stage_type=stage_type, **client_methods)
+    if "collective_rpc_async" not in client_methods:
+        client.collective_rpc_async = AsyncMock(return_value={"via": "collective"})
+    pool = StagePool(0, [client])  # type: ignore[arg-type]
+    return pool, client
+
+
+@pytest.mark.cpu
+def test_collective_rpc_normalizes_none_args_on_control_helper():
+    async def run() -> None:
+        pause = AsyncMock(return_value="paused")
+        pool, client = _make_pool(pause_scheduler_async=pause)
+
+        result = await pool.collective_rpc(0, "pause_scheduler", args=None, kwargs={"mode": "abort"})
+
+        assert result == "paused"
+        pause.assert_awaited_once_with(mode="abort")
+        client.collective_rpc_async.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_unrelated_async_helper_uses_collective_path():
+    async def run() -> None:
+        other = AsyncMock(return_value="should-not-run")
+        pool, client = _make_pool(reset_prefix_cache_async=other)
+
+        result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1})
+
+        assert result == {"via": "collective"}
+        other.assert_not_awaited()
+        client.collective_rpc_async.assert_awaited_once_with(
+            method="reset_prefix_cache",
+            timeout=1.5,
+            args=("x",),
+            kwargs={"k": 1},
+        )
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_control_helper_honors_timeout():
+    async def run() -> None:
+        async def slow_sleep(*_args, **_kwargs):
+            await asyncio.sleep(1.0)
+            return "slept"
+
+        pool, _client = _make_pool(sleep_async=slow_sleep)
+        result = await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,))
+        assert result != "slept"
+        assert result["supported"] is False
+
+    asyncio.run(run())

tests/entrypoints/test_async_omni.py

@@ -185,6 +185,58 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_abort_pops_request_states_only_after_ack():
+    async def run():
+        release = asyncio.Event()
+        seen_during_wait: list[int] = []
+
+        async def slow_abort_async(request_ids):
+            del request_ids
+            seen_during_wait.append(len(omni.request_states))
+            await release.wait()
+
+        omni = get_async_omni_instance(fake_abort_request=slow_abort_async)
+        omni.request_states["req-1-aaaa"] = SimpleNamespace(
+            request_id="req-1-aaaa",
+            external_request_id="req-1",
+            input_stream_task=None,
+        )
+
+        task = asyncio.create_task(omni.abort("req-1"))
+        await asyncio.sleep(0)
+        assert not task.done()
+        assert "req-1-aaaa" in omni.request_states
+        assert seen_during_wait == [1]
+
+        release.set()
+        await task
+        assert "req-1-aaaa" not in omni.request_states
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_propagates_engine_errors_without_popping_state():
+    async def run():
+        async def failing_abort_async(request_ids):
+            del request_ids
+            raise RuntimeError("orchestrator abort failed")
+
+        omni = get_async_omni_instance(fake_abort_request=failing_abort_async)
+        omni.request_states["req-1-bbbb"] = SimpleNamespace(
+            request_id="req-1-bbbb",
+            external_request_id="req-1",
+            input_stream_task=None,
+        )
+
+        with pytest.raises(RuntimeError, match="orchestrator abort failed"):
+            await omni.abort("req-1")
+        assert "req-1-bbbb" in omni.request_states
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_generate_accepts_request_after_repeated_cancellations():
     async def run_test():

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -0,0 +1,310 @@
+"""Unit tests for AR EngineCore vs diffusion worker pause/sleep routing."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, call
+
+import pytest
+
+from vllm_omni.diffusion.data import CuMemTag, OmniACK
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
+    omni = object.__new__(AsyncOmni)
+    omni._pause_cond = asyncio.Condition()
+    omni._paused = False
+    omni._hold_admission_until_resume = False
+    omni._sleeping_tags = set()
+    omni._level2_sleeping = False
+    omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock())
+    omni._final_output_handler = lambda: None
+
+    stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
+    omni.engine = SimpleNamespace(
+        stage_clients=stage_clients,
+        stage_vllm_configs=[None] * len(stage_types),
+        collective_rpc_async=AsyncMock(return_value=[True]),
+    )
+    omni.collective_rpc = AsyncMock(return_value=[True])
+    return omni
+
+
+@pytest.mark.cpu
+def test_split_stage_ids_by_type():
+    omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+    ar_ids, diff_ids = omni._split_stage_ids_by_type()
+    assert ar_ids == [0, 2]
+    assert diff_ids == [1]
+
+
+@pytest.mark.cpu
+def test_split_stage_ids_by_type_rejects_out_of_range():
+    omni = _make_omni(stage_types=["llm", "diffusion"])
+    with pytest.raises(ValueError, match=r"Invalid stage_ids \[2\].*0\.\.1"):
+        omni._split_stage_ids_by_type([0, 2])
+    with pytest.raises(ValueError, match=r"Invalid stage_ids \[-1\]"):
+        omni._split_stage_ids_by_type([-1])
+
+
+@pytest.mark.cpu
+def test_pause_generation_routes_ar_via_collective_rpc():
+    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="abort", clear_cache=True)
+
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        omni.collective_rpc.assert_awaited_once_with(
+            method="pause_scheduler",
+            args=(),
+            kwargs={"mode": "abort", "clear_cache": True},
+            stage_ids=[0],
+        )
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_pause_generation_still_rpcs_when_already_paused():
+    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()
+        omni._paused = True
+
+        await omni.pause_generation(mode="abort", clear_cache=True, stage_ids=[0])
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="pause_scheduler",
+            args=(),
+            kwargs={"mode": "abort", "clear_cache": True},
+            stage_ids=[0],
+        )
+        omni.reset_prefix_cache.assert_awaited_once()
+        omni.reset_mm_cache.assert_awaited_once()
+        omni.reset_encoder_cache.assert_awaited_once()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_resume_generation_resumes_ar_then_clears_frontend_pause():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni._paused = True
+
+        await omni.resume_generation()
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="resume_scheduler",
+            args=(),
+            kwargs=None,
+            stage_ids=[0],
+        )
+        assert omni._paused is False
+        assert omni._hold_admission_until_resume is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
+        omni._sleep_diffusion = AsyncMock(return_value=[diffusion_ack])
+
+        acks = await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="sleep",
+            args=(1, "abort"),
+            kwargs=None,
+            stage_ids=[0],
+        )
+        omni._sleep_diffusion.assert_awaited_once_with([1], 1)
+        assert {ack.stage_id for ack in acks} == {0, 1}
+        assert any(ack.metadata.get("path") == "engine_core" for ack in acks if ack.stage_id == 0)
+        assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+        assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
+        # Sleep gates frontend admission for the trainer resume contract.
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_blocks_admission_before_engine_core_rpc():
+    """Problem 3: _paused must be set before awaiting sleep RPC."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        paused_at_rpc: list[bool] = []
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                paused_at_rpc.append(omni._paused)
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+        assert omni._paused is False
+        await omni.sleep(level=1, mode="abort")
+
+        assert paused_at_rpc == [True]
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        # Sleep must not also call pause_scheduler (EngineCore.sleep pauses).
+        assert all(c.kwargs.get("method") != "pause_scheduler" for c in omni.collective_rpc.await_args_list)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni._paused = True
+        omni._hold_admission_until_resume = True
+        omni._sleeping_tags = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value}
+        diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
+        omni._wake_diffusion = AsyncMock(return_value=[diffusion_ack])
+
+        acks = await omni.wake_up(stage_ids=[0, 1])
+
+        omni.collective_rpc.assert_awaited_once()
+        wake_kwargs = omni.collective_rpc.await_args.kwargs
+        assert wake_kwargs["method"] == "wake_up"
+        assert wake_kwargs["stage_ids"] == [0]
+        assert set(wake_kwargs["kwargs"]["tags"]) == {
+            CuMemTag.WEIGHTS.value,
+            CuMemTag.KV_CACHE.value,
+        }
+        omni._wake_diffusion.assert_awaited_once()
+        assert {ack.stage_id for ack in acks} == {0, 1}
+        assert not omni._sleeping_tags
+        # Mixed/AR wake restores memory but does not resume frontend admission.
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_does_not_resume_frontend_admission():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        await omni.sleep(level=1, mode="abort")
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+        await omni.wake_up()
+
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        assert not omni._sleeping_tags
+        # Explicit resume is required after AR sleep/wake.
+        await omni.resume_generation()
+        assert omni._paused is False
+        assert omni._hold_admission_until_resume is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_level1_wake_without_tags_clears_all_sleeping_tags():
+    """sleep(level=1) tracks WEIGHTS+KV; untagged wake_up must clear both."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+
+        await omni.sleep(level=1, mode="abort")
+        assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+        assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
+
+        await omni.wake_up()
+        assert not omni._sleeping_tags
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_diffusion_only_skips_engine_core_collective_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+
+        await omni.sleep(level=1)
+
+        # AR EngineCore sleep path must not run for diffusion-only.
+        assert call(method="sleep", args=(1, "abort"), kwargs=None, stage_ids=[0]) not in (
+            omni.collective_rpc.await_args_list
+        )
+        omni._sleep_diffusion.assert_awaited_once_with([0], 1)
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_restores_admission_for_diffusion_only():
+    """Pure diffusion must keep sleep → wake → generate (no resume)."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+        omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
+
+        await omni.sleep(level=1)
+        assert omni._paused is True
+
+        async def wait_for_admission() -> None:
+            async with omni._pause_cond:
+                await omni._pause_cond.wait_for(lambda: not omni._paused)
+
+        waiter = asyncio.create_task(wait_for_admission())
+        await omni.wake_up()
+        await asyncio.wait_for(waiter, timeout=1.0)
+        assert omni._paused is False
+        assert omni._hold_admission_until_resume is False
+        assert not omni._sleeping_tags
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion():
+    """Explicit pause_generation still requires resume after diffusion wake."""
+
+    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()
+        omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+        omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
+
+        await omni.pause_generation()
+        await omni.sleep(level=1)
+        await omni.wake_up()
+
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        await omni.resume_generation()
+        assert omni._paused is False
+
+    asyncio.run(run())

tests/entrypoints/test_omni_sleep_mode.py

@@ -114,6 +114,10 @@ async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
         await engine.wake_up(stage_ids=stage_ids)
     except Exception as e:
         logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e)
+    try:
+        await engine.resume_generation(stage_ids=stage_ids)
+    except Exception as e:
+        logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e)
 
 
 def _build_llm_stages() -> tuple[list[dict], list[dict]]:
@@ -240,6 +244,9 @@ class TestOmniLlmSleepMode:
         try:
             await llm_engine.sleep(stage_ids=[0], level=1)
             await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+            # wake_up restores weights but keeps the AR admission hold.
+            # Resume so generate() can reach the sleeping-tags rejection.
+            await llm_engine.resume_generation(stage_ids=[0])
             with pytest.raises(RuntimeError, match="partially or fully asleep"):
                 async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
                     pass
@@ -490,6 +497,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int):
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])
+        await engine.resume_generation(stage_ids=[0, 1])
         async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]):
             pass
     finally:

vllm_omni/engine/async_omni_engine.py

@@ -50,6 +50,7 @@ from vllm_omni.engine.async_engine_utils import (
 )
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -1736,14 +1737,51 @@ class AsyncOmniEngine:
         return self.stage_metadata[stage_id]
 
     def abort(self, request_ids: list[str]) -> None:
-        """Send abort message to the Orchestrator."""
+        """Fire-and-forget abort: enqueue and return without waiting.
+
+        Prefer :meth:`abort_async` when the caller needs acknowledgment that
+        stage aborts, binding release, and orchestrator request cleanup finished.
+        """
         if self.request_queue is None:
             raise RuntimeError("request_queue is not initialized")
         self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids))
 
-    async def abort_async(self, request_ids: list[str]) -> None:
-        """Async abort API."""
-        self.abort(request_ids)
+    async def abort_async(
+        self,
+        request_ids: list[str],
+        timeout: float | None = None,
+    ) -> None:
+        """Abort requests and wait for orchestrator acknowledgment.
+
+        Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the
+        :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and
+        raises if the orchestrator reports failure or times out.
+        """
+        if self.request_queue is None:
+            raise RuntimeError("request_queue is not initialized")
+        transport = self._correlated_rpc_client
+        if transport is None:
+            raise RuntimeError("correlated RPC client is not initialized")
+
+        rpc_id = uuid.uuid4().hex
+        msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id)
+
+        def _wait() -> AbortResultMessage:
+            result_msg = transport.execute(
+                ("abort", rpc_id),
+                msg,
+                timeout=timeout,
+                timeout_message=f"abort timed out after {timeout} seconds",
+                block_on_submit=True,
+            )
+            if not isinstance(result_msg, AbortResultMessage):
+                raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}")
+            return result_msg
+
+        loop = asyncio.get_running_loop()
+        result_msg = await loop.run_in_executor(None, _wait)
+        if not result_msg.success:
+            raise RuntimeError(result_msg.error or "abort failed")
 
     def submit_interaction(
         self,

vllm_omni/engine/messages.py

@@ -42,6 +42,20 @@ class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True):
 class AbortRequestMessage(EngineQueueMessage, kw_only=True):
     type: Literal["abort"] = "abort"
     request_ids: list[str]
+    # When set, Orchestrator emits AbortResultMessage on rpc_async_queue so
+    # callers can await acknowledgment. Omit for fire-and-forget abort().
+    rpc_id: str | None = None
+
+
+class AbortResultMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["abort_result"] = "abort_result"
+    rpc_id: str
+    success: bool
+    error: str | None = None
+
+    @property
+    def rpc_correlation_key(self) -> tuple[str, str]:
+        return ("abort", self.rpc_id)
 
 
 class InteractionMessage(EngineQueueMessage, kw_only=True):

vllm_omni/engine/orchestrator.py

@@ -36,6 +36,7 @@ from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.engine.membership_controller import MembershipController
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -817,13 +818,35 @@ class Orchestrator:
         )
 
     async def _handle_abort(self, msg: AbortRequestMessage) -> None:
-        """Handle an abort message from the main thread."""
+        """Handle an abort message from the main thread.
+
+        When ``msg.rpc_id`` is set, emit :class:`AbortResultMessage` after
+        stage aborts, binding release, and request cleanup complete so the
+        caller can await acknowledgment. Fire-and-forget aborts omit rpc_id.
+        """
         request_ids = msg.request_ids
-        # _cleanup_request_ids is CFG-aware: it expands aborted parents to
-        # their companions and fails a deferred parent whose companion is
-        # aborted before its output arrived.
-        await self._cleanup_request_ids(list(request_ids), abort=True)
-        logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
+        error: str | None = None
+        try:
+            # _cleanup_request_ids is CFG-aware: it expands aborted parents to
+            # their companions and fails a deferred parent whose companion is
+            # aborted before its output arrived.
+            await self._cleanup_request_ids(list(request_ids), abort=True)
+            logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
+        except Exception as exc:
+            error = str(exc)
+            logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids)
+            if msg.rpc_id is None:
+                raise
+        if msg.rpc_id is not None:
+            # Always emit a result when rpc_id is set so CorrelatedRpcClient
+            # waiters are unblocked even on failure.
+            await self.rpc_async_queue.put(
+                AbortResultMessage(
+                    rpc_id=msg.rpc_id,
+                    success=error is None,
+                    error=error,
+                )
+            )
 
     async def _handle_interaction(self, msg: InteractionMessage) -> None:
         """Handle a midway interaction for an active streaming diffusion request."""

vllm_omni/engine/stage_pool.py

@@ -94,6 +94,16 @@ class StagePool:
 
     DISPATCH_WAIT_TIMEOUT_S: float = 10.0
     DISPATCH_RETRY_INTERVAL_S: float = 0.1
+    # Only these EngineCore helpers may skip collective_rpc_async. A generic
+    # ``{method}_async`` on AsyncMPClient must not silently drop timeout.
+    _ENGINE_CORE_CONTROL_ASYNC_METHODS = frozenset(
+        {
+            "pause_scheduler",
+            "resume_scheduler",
+            "sleep",
+            "wake_up",
+        }
+    )
 
     def __init__(
         self,
@@ -1214,6 +1224,7 @@ class StagePool:
         kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any] | Any:
         """Dispatch a stage-scoped control-plane RPC to one physical route."""
+        args = tuple(args or ())
         kwargs = dict(kwargs or {})
         client = self.clients[replica_id]
         if client is None:
@@ -1222,6 +1233,14 @@ class StagePool:
                 "error": f"stage {self.stage_id} replica {replica_id} is not attached",
             }
         try:
+            if self.stage_type != "diffusion" and method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS:
+                client_method = getattr(client, f"{method}_async", None)
+                if callable(client_method):
+                    result = client_method(*args, **kwargs)
+                    if timeout is not None:
+                        return await asyncio.wait_for(result, timeout=timeout)
+                    return await result
+
             return await client.collective_rpc_async(
                 method=method,
                 timeout=timeout,

vllm_omni/entrypoints/async_omni.py

@@ -146,6 +146,11 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
+        # True after pause_generation() or AR EngineCore sleep; wake_up must
+        # not reopen generate() until resume_generation(). Diffusion-only
+        # 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
         self._sleeping_tags: set[str] = set()
         self._level2_sleeping: bool = False
         self._duplex_request_client: DuplexRequestClient | None = None
@@ -1009,6 +1014,30 @@ class AsyncOmni(EngineClient, OmniBase):
 
         return results
 
+    async def _engine_core_rpc(
+        self,
+        method: str,
+        *,
+        stage_ids: list[int],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """Call an AR EngineCore 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.
+        """
+        results = await self.collective_rpc(
+            method=method,
+            args=args,
+            kwargs=kwargs,
+            stage_ids=stage_ids,
+        )
+        for result in results:
+            if isinstance(result, dict) and result.get("error"):
+                raise RuntimeError(f"{method} failed: {result['error']}")
+        return results
+
     @staticmethod
     def _coerce_stage_bool(result: Any) -> bool:
         """Reduce a stage RPC result to a boolean.
@@ -1079,7 +1108,12 @@ class AsyncOmni(EngineClient, OmniBase):
         await self._abort(internal_req_ids)
 
     async def _abort(self, request_ids: list[str]) -> None:
-        """Submit request IDs to be aborted to the engine."""
+        """Abort request IDs via the engine and clean frontend state after ack.
+
+        Waits for orchestrator abort acknowledgment before popping
+        ``request_states`` so generate() cleanup stays consistent with
+        backend binding/request teardown. Orchestrator abort errors propagate.
+        """
         await self.engine.abort_async(request_ids)
         for rid in request_ids:
             state = self.request_states.pop(rid, None)
@@ -1089,23 +1123,77 @@ class AsyncOmni(EngineClient, OmniBase):
         if self.log_stats:
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 
+    def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]:
+        """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC)."""
+        n_stages = len(self.engine.stage_clients)
+        if stage_ids is None:
+            stage_ids = list(range(n_stages))
+        else:
+            invalid = [sid for sid in stage_ids if not isinstance(sid, int) or sid < 0 or sid >= n_stages]
+            if invalid:
+                raise ValueError(
+                    f"Invalid stage_ids {invalid}; valid range is 0..{n_stages - 1}"
+                    if n_stages
+                    else f"Invalid stage_ids {invalid}; this engine has no stages"
+                )
+        ar_stage_ids: list[int] = []
+        diffusion_stage_ids: list[int] = []
+        for sid in stage_ids:
+            client = self.engine.stage_clients[sid]
+            if getattr(client, "stage_type", "llm") == "diffusion":
+                diffusion_stage_ids.append(sid)
+            else:
+                ar_stage_ids.append(sid)
+        return ar_stage_ids, diffusion_stage_ids
+
     async def pause_generation(
         self,
         *,
         mode: PauseMode = "abort",
         wait_for_inflight_requests: bool = False,
         clear_cache: bool = True,
+        stage_ids: list[int] | None = None,
     ) -> None:
-        """Pause generation."""
+        """Pause generation, mirroring vLLM AsyncLLM.pause_generation.
+
+        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.
+
+        Note: ``sleep()`` already pauses the AR scheduler internally (same as
+        vLLM EngineCore.sleep). Call this API when you need pause *without*
+        freeing GPU memory (e.g. weight sync).
+        """
+        if wait_for_inflight_requests:
+            mode = "wait"
+
         async with self._pause_cond:
-            if self._paused:
-                return
+            # Keep running EngineCore pause + cache clear even when frontend
+            # admission is already paused (sleep or a prior pause_generation).
             self._paused = True
+            self._hold_admission_until_resume = True
 
-        # TODO: Implement request draining if wait_for_inflight_requests
+        ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        if ar_stage_ids:
+            logger.info(
+                "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)",
+                self._name,
+                ar_stage_ids,
+                mode,
+            )
+            # Same API name as vLLM AsyncMPClient.pause_scheduler_async; routed
+            # through collective_rpc so it runs on the orchestrator event loop.
+            await self._engine_core_rpc(
+                "pause_scheduler",
+                stage_ids=ar_stage_ids,
+                kwargs={"mode": mode, "clear_cache": clear_cache},
+            )
 
+        # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler
+        # already clears AR-side caches when clear_cache=True.
         if clear_cache:
-            # Clear caches for all stages.
             await self.reset_prefix_cache(
                 reset_running_requests=not wait_for_inflight_requests,
                 reset_connector=True,
@@ -1113,14 +1201,20 @@ class AsyncOmni(EngineClient, OmniBase):
             await self.reset_mm_cache()
             await self.reset_encoder_cache()
 
-    async def resume_generation(self) -> None:
-        """Resume generation."""
+    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)
+        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)
+
         async with self._pause_cond:
             self._paused = False
+            self._hold_admission_until_resume = False
             self._pause_cond.notify_all()
 
     async def is_paused(self) -> bool:
-        """Check if paused."""
+        """Check if frontend admission is paused."""
         async with self._pause_cond:
             return self._paused
 
@@ -1178,33 +1272,65 @@ class AsyncOmni(EngineClient, OmniBase):
     async def sleep(
         self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort"
     ) -> list[OmniACK]:
+        """Put stages to sleep.
+
+        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.
+
+        Frontend admission is blocked at the start of this call (``_paused``)
+        so pipelined :meth:`generate` cannot race into stages while sleep is
+        in flight. This does **not** invoke EngineCore.pause_scheduler again
+        (sleep already pauses the AR scheduler).
+
+        For AR / mixed engines, ``wake_up`` does **not** clear ``_paused``;
+        callers must :meth:`resume_generation` when ready (typical trainer
+        order: pause → abort → sleep → train → wake → resume). Diffusion-only
+        engines have no EngineCore pause to hold, so ``wake_up`` restores
+        admission and ``sleep → wake → generate`` keeps working.
+        """
+        # Block admission before any sleep RPC so generate() waits on
+        # _pause_cond during the drain/offload window. EngineCore.sleep will
+        # pause the AR scheduler itself; avoid a second pause_scheduler here.
+        async with self._pause_cond:
+            self._paused = True
+
         self._final_output_handler()
-        if stage_ids is None:
-            stage_ids = list(range(len(self.engine.stage_clients)))
-        total_workers = 0
-        for sid in stage_ids:
-            client = self.engine.stage_clients[sid]
-            # During the Diffusion phase, regardless of the TP amount,
-            # currently only a summary ACK is reported at Rank 0.
-            if getattr(client, "stage_type", "") == "diffusion":
-                total_workers += 1
-            else:
-                config = self.engine.stage_vllm_configs[sid]
-                actual_tp = config.parallel_config.tensor_parallel_size if config else 1
-                total_workers += actual_tp
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        final_acks: list[OmniACK] = []
+        if ar_stage_ids:
+            self._hold_admission_until_resume = True
+            logger.info(
+                "[%s] Sleeping AR stage(s) %s via EngineCore.sleep(level=%s, mode=%s)",
+                self._name,
+                ar_stage_ids,
+                level,
+                mode,
+            )
+            await self._engine_core_rpc(
+                "sleep",
+                stage_ids=ar_stage_ids,
+                args=(level, mode),
+            )
+            # EngineCore.sleep has no OmniACK handshake; emit stage-level SUCCESS
+            # markers so callers/tests that count ACKs keep a stable API.
+            task_id = f"engine_core-sleep-{uuid.uuid4().hex[:8]}"
+            final_acks.extend(
+                OmniACK(
+                    task_id=task_id,
+                    status="SUCCESS",
+                    stage_id=sid,
+                    rank=0,
+                    metadata={"path": "engine_core", "level": level, "mode": mode},
+                )
+                for sid in ar_stage_ids
+            )
+
+        if diffusion_stage_ids:
+            final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level))
 
-        task_id = str(uuid.uuid4())
-        self.event_resolver.watch_task(task_id, expected_count=total_workers)
-        logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniSleepTask(level=level, task_id=task_id)
-        rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
-        final_acks = []
-        for stage_res in rpc_results:
-            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
-            for ack in worker_acks:
-                if ack is not None:
-                    await self.event_resolver.resolve(ack)
-                    final_acks.append(ack)
         if not hasattr(self, "_sleeping_tags"):
             self._sleeping_tags = set()
         self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value])
@@ -1212,7 +1338,36 @@ class AsyncOmni(EngineClient, OmniBase):
             self._level2_sleeping = True
         return final_acks
 
+    async def _sleep_diffusion(self, stage_ids: list[int], level: int) -> list[OmniACK]:
+        """Worker-level sleep RPC for diffusion stages only."""
+        # Diffusion reports one summary ACK at rank 0 regardless of TP.
+        total_workers = len(stage_ids)
+        task_id = str(uuid.uuid4())
+        self.event_resolver.watch_task(task_id, expected_count=total_workers)
+        logger.info("[%s] Sleep (diffusion) initiated (Task: %s).", self._name, task_id)
+        task = OmniSleepTask(level=level, task_id=task_id)
+        rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
+        final_acks: list[OmniACK] = []
+        for stage_res in rpc_results:
+            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
+            for ack in worker_acks:
+                if ack is not None:
+                    await self.event_resolver.resolve(ack)
+                    final_acks.append(ack)
+        return final_acks
+
     async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
+        """Wake stages after sleep.
+
+        AR/LLM stages use EngineCore.wake_up (restore memory, auto-resume
+        scheduler). Diffusion stages keep the worker-level wake RPC.
+
+        Does **not** clear the frontend ``_paused`` admission gate when
+        :meth:`pause_generation` ran or AR stages were slept — call
+        :meth:`resume_generation` when the trainer is ready to admit new
+        requests. Diffusion-only ``sleep`` uses ``_paused`` only as a race
+        guard; this method restores admission after a successful wake.
+        """
         self._final_output_handler()
 
         if getattr(self, "_level2_sleeping", False):
@@ -1231,30 +1386,30 @@ class AsyncOmni(EngineClient, OmniBase):
             logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
             return []
 
-        if stage_ids is None:
-            stage_ids = list(range(len(self.engine.stage_clients)))
-        total_workers = 0
-        for sid in stage_ids:
-            client = self.engine.stage_clients[sid]
-            if getattr(client, "stage_type", "") == "diffusion":
-                total_workers += 1
-            else:
-                config = self.engine.stage_vllm_configs[sid]
-                total_workers += config.parallel_config.tensor_parallel_size if config else 1
-        task_id = str(uuid.uuid4())
-        self.event_resolver.watch_task(task_id, expected_count=total_workers)
-        logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
-        rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
-        final_acks = []
-        for stage_res in rpc_results:
-            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
-            for ack in worker_acks:
-                if ack is not None:
-                    await self.event_resolver.resolve(ack)
-                    final_acks.append(ack)
-        current_omni_platform.synchronize()
-        await asyncio.sleep(0.1)
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+
+        final_acks: list[OmniACK] = []
+        if ar_stage_ids:
+            logger.info("[%s] Waking AR stage(s) %s via EngineCore", self._name, ar_stage_ids)
+            await self._engine_core_rpc(
+                "wake_up",
+                stage_ids=ar_stage_ids,
+                kwargs={"tags": requested_tags},
+            )
+            task_id = f"engine_core-wake-{uuid.uuid4().hex[:8]}"
+            final_acks.extend(
+                OmniACK(
+                    task_id=task_id,
+                    status="SUCCESS",
+                    stage_id=sid,
+                    rank=0,
+                    metadata={"path": "engine_core", "tags": list(requested_tags)},
+                )
+                for sid in ar_stage_ids
+            )
+
+        if diffusion_stage_ids:
+            final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags))
 
         for t in requested_tags:
             if hasattr(self, "_sleeping_tags"):
@@ -1263,7 +1418,37 @@ class AsyncOmni(EngineClient, OmniBase):
         # wake support (e.g. tags=["kv_cache"] only) is added in the future.
         if not getattr(self, "_sleeping_tags", None):
             self._level2_sleeping = False
-        logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
+        logger.info(
+            "[%s] Wake-up complete for stage(s) %s.",
+            self._name,
+            ar_stage_ids + diffusion_stage_ids,
+        )
+        # Diffusion-only sleep uses `_paused` as a race guard. Restore
+        # generate() admission after memory is back. AR/mixed sleep and
+        # pause_generation keep the trainer hold until resume_generation.
+        if not getattr(self, "_hold_admission_until_resume", False):
+            async with self._pause_cond:
+                self._paused = False
+                self._pause_cond.notify_all()
+        return final_acks
+
+    async def _wake_diffusion(self, stage_ids: list[int], requested_tags: list[str]) -> list[OmniACK]:
+        """Worker-level wake RPC for diffusion stages only."""
+        total_workers = len(stage_ids)
+        task_id = str(uuid.uuid4())
+        self.event_resolver.watch_task(task_id, expected_count=total_workers)
+        logger.info("[%s] Wake-up (diffusion) initiated (Task: %s).", self._name, task_id)
+        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
+        rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
+        final_acks: list[OmniACK] = []
+        for stage_res in rpc_results:
+            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
+            for ack in worker_acks:
+                if ack is not None:
+                    await self.event_resolver.resolve(ack)
+                    final_acks.append(ack)
+        current_omni_platform.synchronize()
+        await asyncio.sleep(0.1)
         return final_acks
 
     async def is_sleeping(self) -> bool: