[Bugfix] Fix colocate-async sleep admission race and deliver AR abort tokens (#6367)

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

.buildkite/amd/test-amd-merge.yml

@@ -175,7 +175,7 @@ steps:
       grade: Blocking
       commands:
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
 
     - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
       agent_pool: mi300_1

.buildkite/cuda/test-merge.yml

@@ -122,6 +122,7 @@ steps:
       - label: "Omni · Qwen3-Omni Test"
         source_file_dependencies:
           - tests/e2e/offline_inference/test_qwen3_omni.py
+          - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py
           - tests/e2e/online_serving/test_qwen3_omni.py
           - tests/helpers/runtime.py
           - vllm_omni/entrypoints/openai/serving_chat.py
@@ -132,9 +133,10 @@ steps:
           - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
           - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
           - vllm_omni/deploy/qwen3_omni_moe.yaml
+          - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
         timeout_in_minutes: 50
         commands:
-          - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_2
 
       - label: "Omni · MiniCPM-o 4.5 Test"

docs/design/module/engine_orchestration.md

@@ -37,6 +37,7 @@ validation_paths:
   - tests/engine/test_async_omni_engine_input.py
   - tests/engine/test_async_omni_engine_outputs.py
   - tests/engine/test_async_omni_engine_abort.py
+  - tests/e2e/offline_inference/test_qwen3_omni.py
   - tests/engine/test_async_omni_engine_stage_init.py
   - tests/engine/test_orchestrator.py
   - tests/engine/test_orchestrator_error_handling.py
@@ -113,6 +114,15 @@ Numbers become append-only after normative promotion.
 
 Test routing, output ordering, queue correlation, cancellation, failure
 propagation, shutdown ordering, and representative multi-stage execution.
+AR abort of a final-stage LLM request must deliver a terminal output with
+`finish_reason="abort"` and the generated prefix so collocated training can
+resume; diffusion abort remains whole-sample retry. Aborting a parallel-sampling
+child must drop `parent_requests` once no children remain. EngineCore control
+RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must
+propagate worker exceptions rather than returning `{"supported": False}`.
+When a request has multiple final output stages, only the last abort message
+is request-terminal. Output-processor abort state is committed only after the
+physical EngineCore abort succeeds.
 Changes that cross into stage runtime or public error behavior require review
 from that contract's owners.
 

docs/design/module/entrypoints.md

@@ -43,6 +43,7 @@ depends_on:
 validation_paths:
   - tests/entrypoints/test_omni_entrypoints.py
   - tests/entrypoints/test_async_omni.py
+  - tests/entrypoints/test_async_omni_pause_sleep_routing.py
   - tests/entrypoints/test_async_omni_duplex.py
   - tests/entrypoints/test_serve.py
   - tests/entrypoints/test_stream_finish_reason.py
@@ -119,7 +120,13 @@ promotion.
 
 Test request validation, protocol conversion, model-adapter routing,
 streaming, session identity, disconnect/cancellation, and error mapping for
-each affected offline or serving entrypoint.
+each affected offline or serving entrypoint. Sleep must wait for in-flight
+`generate()` admission before EngineCore offload; `wake_up` does not resume
+admission — callers must `resume_generation()`. Sleeping tags are tracked per
+stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`.
+Streaming input pumps take an admission slot immediately before each EngineCore
+ADD or update, not while waiting for the next client chunk. Frontend abort
+keeps `request_states` until `generate()` consumes the terminal output.
 
 ## Promotion gate
 

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -3,6 +3,9 @@
 
 """
 E2E offline tests for Omni model with video input and audio output.
+
+Abort / sleep-admission lives in ``test_qwen3_omni_colocate_async.py`` so it
+does not overlap this module's OmniRunners.
 """
 
 import os

tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Isolated AsyncOmni abort / sleep-admission regression for Qwen3-Omni.
+
+This lives in its own module so the L3 merge job can run it *after*
+``test_qwen3_omni.py`` (offline) and the online serving file. Those modules keep
+module-scoped OmniRunner / OmniServer instances until teardown; overlapping a
+CuMem sleep-mode engine with them kills StageEngineCoreProc at handshake
+(``Failed core proc(s): {}``).
+"""
+
+import asyncio
+import os
+from contextlib import ExitStack
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+from vllm import SamplingParams
+from vllm.inputs import TokensPrompt
+
+from tests.helpers.clean import wait_for_gpu_memory_to_clear
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.platforms import current_omni_platform
+
+MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
+
+def _output_token_ids(output) -> list[int]:
+    if not getattr(output, "outputs", None):
+        return []
+    completion = output.outputs[0]
+    token_ids = list(completion.token_ids or [])
+    if token_ids:
+        return token_ids
+    return list(getattr(completion, "cumulative_token_ids", None) or [])
+
+
+def _finish_reason(output) -> str | None:
+    if not getattr(output, "outputs", None):
+        return None
+    return getattr(output.outputs[0], "finish_reason", None)
+
+
+def _colocate_async_deploy() -> str:
+    """Thinker-only Instruct deploy with sleep mode on a single GPU.
+
+    Abort / sleep admission is a thinker control-plane contract. The 3-stage
+    CI yaml plus CuMem on every stage fails engine-core startup in the L3
+    merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("qwen3_omni_moe_thinking.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enable_sleep_mode": True,
+                    "enforce_eager": True,
+                    "trust_remote_code": True,
+                    "tensor_parallel_size": 1,
+                    "devices": "0",
+                    "gpu_memory_utilization": 0.9,
+                    "max_num_seqs": 1,
+                    "max_model_len": 8192,
+                    "max_num_batched_tokens": 8192,
+                    "enable_prefix_caching": False,
+                    "skip_mm_profiling": True,
+                },
+            },
+        },
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+@pytest.mark.asyncio
+async def test_colocate_async_abort_tokens_and_sleep_admission() -> None:
+    """**required** regression. Control-plane APIs (``abort`` / ``sleep`` /
+    ``wake_up`` / ``resume_generation``) are not in ``send_omni_request``.
+
+    On ``origin/main`` this test fails:
+    - abort: ``generate()`` never yields ``finish_reason="abort"`` with the
+      tokens produced so far, so resume has an empty prefix (or the stream
+      hangs after frontend state is dropped).
+    - sleep: ``generate()`` is admitted into EngineCore while weights/KV are
+      offloaded, so the task errors (asleep / corrupted ADD frame) instead of
+      waiting for ``resume_generation()``.
+
+    On this PR both legs pass: abort returns a resumeable prefix, and sleep
+    holds ``generate()`` until wake + resume.
+    """
+    prompt = "What color is the sky? Write a long, detailed explanation."
+    device_count = current_omni_platform.device_count()
+    if device_count > 0:
+        # Fail closed if a prior module-scoped runner is still on the GPU.
+        # ``clean_gpu_memory_between_tests`` only logs a note on timeout.
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(device_count)),
+            threshold_ratio=0.15,
+            timeout_s=120,
+        )
+
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            deploy_config=_colocate_async_deploy(),
+            enable_sleep_mode=True,
+        )
+        after.callback(engine.shutdown)
+
+        request_id = "qwen3-abort-partial"
+        outputs: list = []
+
+        async def _generate(max_tokens: int, req_id: str, gen_prompt) -> None:
+            async for output in engine.generate(
+                prompt=gen_prompt,
+                request_id=req_id,
+                sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens),
+                output_modalities=["text"],
+            ):
+                outputs.append(output)
+
+        abort_task = asyncio.create_task(_generate(256, request_id, prompt))
+        prefix: list[int] = []
+        prompt_token_ids: list[int] = []
+        for _ in range(600):
+            if abort_task.done():
+                break
+            if outputs:
+                latest = outputs[-1]
+                prefix = _output_token_ids(latest)
+                prompt_token_ids = list(getattr(latest, "prompt_token_ids", None) or [])
+                if prefix:
+                    break
+            await asyncio.sleep(0.1)
+
+        assert prefix, "generate produced no tokens before abort"
+        assert not abort_task.done(), "generate finished before abort; raise max_tokens"
+        await engine.abort(request_id)
+        await asyncio.wait_for(abort_task, timeout=90)
+
+        final = outputs[-1]
+        assert final.finished
+        assert _finish_reason(final) == "abort"
+        abort_prefix = _output_token_ids(final)
+        assert abort_prefix, "abort dropped the generated prefix (main-branch behavior)"
+
+        resume_prompt: str | TokensPrompt = (
+            TokensPrompt(prompt_token_ids=prompt_token_ids + abort_prefix) if prompt_token_ids else prompt
+        )
+        outputs.clear()
+        async for output in engine.generate(
+            prompt=resume_prompt,
+            request_id="qwen3-abort-resume",
+            sampling_params=SamplingParams(temperature=0.0, max_tokens=8),
+            output_modalities=["text"],
+        ):
+            outputs.append(output)
+        assert outputs
+        assert any(_output_token_ids(out) for out in outputs)
+
+        # Trainer order is pause → abort → sleep. Pause here after the
+        # resume generate so EngineCore is idle before CuMem offload.
+        await engine.pause_generation(mode="abort", clear_cache=True)
+        await engine.sleep(level=1)
+        outputs.clear()
+        sleep_task = asyncio.create_task(_generate(8, "qwen3-sleep-admission", prompt))
+        await asyncio.sleep(1.0)
+        assert not sleep_task.done(), "generate() ran while EngineCore was sleeping (main-branch admission race)"
+        await engine.wake_up()
+        await asyncio.sleep(0.5)
+        assert not sleep_task.done(), "generate() resumed before resume_generation()"
+        await engine.resume_generation()
+        await asyncio.wait_for(sleep_task, timeout=180)
+        assert outputs
+        assert _finish_reason(outputs[-1]) != "abort"
+    await asyncio.sleep(5)

tests/engine/test_async_omni_engine_abort_ack.py

@@ -67,7 +67,8 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
         assert not task.done()
 
         rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
-        await task
+        result = await task
+        assert result == []
 
     try:
         asyncio.run(_run())

tests/engine/test_orchestrator.py

@@ -256,8 +256,32 @@ class FakeOutputProcessor:
         )
 
     def abort_requests(self, request_ids, internal: bool = False):
-        self.abort_calls.append(request_ids)
-        return request_ids
+        aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal)
+        return aborted_ids
+
+    def abort_requests_collecting_outputs(self, request_ids, *, internal: bool = False, commit_state: bool = True):
+        """Mirror MultimodalOutputProcessor collecting for AR abort-prefix tests."""
+        del internal, commit_state
+        ids = list(request_ids)
+        self.abort_calls.append(ids)
+        outputs: list[RequestOutput] = []
+        for rid in ids:
+            seeded = next(
+                (ro for ro in self.request_outputs if getattr(ro, "request_id", None) == rid),
+                None,
+            )
+            token_ids = list(seeded.outputs[0].token_ids) if seeded is not None and seeded.outputs else [1, 2]
+            # Intentionally stamp an internal-looking id on the RequestOutput so
+            # StagePool must re-key by the orchestrator id it passed in.
+            outputs.append(
+                _build_request_output(
+                    f"engine-internal-{rid}",
+                    token_ids=token_ids,
+                    finished=True,
+                    finish_reason="abort",
+                )
+            )
+        return ids, outputs
 
     def update_scheduler_stats(self, _scheduler_stats) -> None:
         return None
@@ -320,6 +344,7 @@ def _build_request_output(
     prompt_token_ids: list[int] | None = None,
     finished: bool = True,
     text: str = "test",
+    finish_reason: str | None = None,
 ) -> RequestOutput:
     completion = CompletionOutput(
         index=0,
@@ -327,7 +352,7 @@ def _build_request_output(
         token_ids=list(token_ids or [1, 2]),
         cumulative_logprob=0.0,
         logprobs=None,
-        finish_reason="stop" if finished else None,
+        finish_reason=(finish_reason if finish_reason is not None else ("stop" if finished else None)),
         stop_reason=None,
     )
     return RequestOutput(
@@ -983,7 +1008,11 @@ async def test_run_abort(orchestrator_factory) -> None:
 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),
+        FakeStageClient(
+            stage_type="llm",
+            final_output=True,
+            next_inputs=[{"prompt_token_ids": [7, 8, 9]}],
+        ),
     ]
     processors = [
         FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]),
@@ -1002,6 +1031,11 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N
             final_stage_id=1,
         )
         await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+        # Abort outputs come from the final AR stage's output processor, and
+        # StagePool.abort_requests only collects for requests with a live
+        # replica binding. Forward off stage-0 first so stage-1 is bound.
+        stages[0].push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+        await _wait_for(lambda: len(stages[1].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)
@@ -1011,11 +1045,53 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N
         assert result.error is None
         assert result.rpc_correlation_key == ("abort", "abort-1")
         assert stages[0].abort_calls == [["req-ack"]]
+        assert stages[1].abort_calls == [["req-ack"]]
         assert "req-ack" not in orchestrator_fixture.orchestrator.request_states
+        assert result.abort_outputs is not None
+        assert len(result.abort_outputs) == 1
+        abort_msg = result.abort_outputs[0]
+        assert abort_msg.request_id == "req-ack"
+        assert abort_msg.finished is True
+        assert abort_msg.stage_id == 1
+        assert list(abort_msg.engine_outputs.outputs[0].token_ids) == [2]
+        assert abort_msg.engine_outputs.outputs[0].finish_reason == "abort"
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
 
 
+@pytest.mark.asyncio
+async def test_abort_marks_only_last_final_stage_output_finished() -> None:
+    """A request with two final AR stages must keep earlier abort outputs unfinished."""
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=True),
+        FakeStageClient(stage_type="llm", final_output=True),
+    ]
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[2], finished=True)]),
+    ]
+    pools = _build_stage_pools([[stages[0]], [stages[1]]], output_processors=processors)
+    pools[0]._request_bindings["req-two"] = 0
+    pools[1]._request_bindings["req-two"] = 0
+
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=pools,
+    )
+    orchestrator.request_states["req-two"] = OrchestratorRequestState(
+        request_id="req-two",
+        final_stage_id=1,
+        final_output_stage_ids={0, 1},
+    )
+
+    outputs = await orchestrator._abort_request_ids(["req-two"])
+    assert [(msg.stage_id, msg.finished) for msg in outputs] == [(0, False), (1, True)]
+    assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [1]
+    assert list(outputs[1].engine_outputs.outputs[0].token_ids) == [2]
+
+
 @pytest.mark.asyncio
 async def test_handle_abort_emits_error_result_on_failure() -> None:
     rpc_queue: asyncio.Queue = asyncio.Queue()

tests/engine/test_output_processor.py

@@ -700,3 +700,62 @@ def test_num_cache_creation_tokens_reaches_direct_request_output():
 
     assert direct.num_cached_tokens == 3
     assert direct.num_cache_creation_tokens == 2
+
+
+def _abort_processor_with_parent(child_ids: list[str]):
+    processor = object.__new__(MultimodalOutputProcessor)
+    processor.request_states = {}
+    processor.external_req_ids = {}
+    processor.parent_requests = {}
+    processor._native_text_metrics_by_request = {}
+    processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None)
+    parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids))
+    processor.parent_requests["parent"] = parent
+    for child_id in child_ids:
+        req_state = SimpleNamespace(
+            parent_req=parent,
+            lora_name=None,
+            output_kind=RequestOutputKind.CUMULATIVE,
+            detokenizer=SimpleNamespace(output_token_ids=[7, 8]),
+            queue=None,
+            external_req_id="parent",
+            make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)),
+        )
+        processor.request_states[child_id] = req_state
+    return processor, parent
+
+
+def test_abort_last_parallel_child_drops_parent_request():
+    processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"])
+
+    aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True)
+    assert aborted == ["0_parent"]
+    assert "parent" in processor.parent_requests
+    assert parent.child_requests == {"1_parent"}
+
+    aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True)
+    assert aborted == ["1_parent"]
+    assert "parent" not in processor.parent_requests
+    assert not parent.child_requests
+
+
+def test_abort_snapshot_leaves_state_until_commit():
+    processor, parent = _abort_processor_with_parent(["0_parent"])
+    processor.external_req_ids["parent"] = ["0_parent"]
+
+    aborted, outputs = processor.abort_requests_collecting_outputs(
+        ["parent"],
+        internal=False,
+        commit_state=False,
+    )
+    assert aborted == ["0_parent"]
+    assert outputs
+    assert "0_parent" in processor.request_states
+    assert processor.external_req_ids["parent"] == ["0_parent"]
+    assert "parent" in processor.parent_requests
+    assert parent.child_requests == {"0_parent"}
+
+    processor.commit_aborted_request_state(["parent"], internal=False)
+    assert "0_parent" not in processor.request_states
+    assert "parent" not in processor.parent_requests
+    assert "parent" not in processor.external_req_ids

tests/engine/test_stage_pool_collective_rpc.py

@@ -64,8 +64,61 @@ def test_collective_rpc_control_helper_honors_timeout():
             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"
+        with pytest.raises(asyncio.TimeoutError):
+            await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,))
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_control_method_reraises_worker_error():
+    async def run() -> None:
+        pool, _client = _make_pool(wake_up_async=AsyncMock(side_effect=RuntimeError("worker died")))
+        with pytest.raises(RuntimeError, match="worker died"):
+            await pool.collective_rpc(0, "wake_up")
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_non_control_method_still_returns_error_dict():
+    async def run() -> None:
+        pool, _client = _make_pool(collective_rpc_async=AsyncMock(side_effect=RuntimeError("probe failed")))
+        result = await pool.collective_rpc(0, "reset_prefix_cache")
         assert result["supported"] is False
+        assert "probe failed" in result["error"]
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
+    async def run() -> None:
+        class RecordingOutputProcessor:
+            def __init__(self) -> None:
+                self.collected = False
+                self.committed = False
+
+            def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True):
+                del internal
+                self.collected = True
+                assert commit_state is False
+                return list(request_ids), [SimpleNamespace(request_id=request_ids[0])]
+
+            def commit_aborted_request_state(self, request_ids, *, internal=False):
+                del request_ids, internal
+                self.committed = True
+
+        abort = AsyncMock(side_effect=RuntimeError("engine abort failed"))
+        output_processor = RecordingOutputProcessor()
+        client = SimpleNamespace(stage_type="llm", abort_requests_async=abort)
+        pool = StagePool(0, [client], output_processor=output_processor)  # type: ignore[arg-type]
+        pool._request_bindings["req-1"] = 0
+
+        with pytest.raises(RuntimeError, match="engine abort failed"):
+            await pool.abort_requests(["req-1"])
+        assert output_processor.collected is True
+        assert output_processor.committed is False
+        abort.assert_awaited_once()
 
     asyncio.run(run())

tests/entrypoints/test_async_omni.py

@@ -180,7 +180,6 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
             assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
         assert len(aborted_ids) == expected_cancel_count
         assert len(set(aborted_ids)) == expected_cancel_count
-        assert len(omni.request_states) == len(req_ids) - expected_cancel_count
 
         for t in tasks:
             t.cancel()
@@ -190,7 +189,9 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
 
 
 @pytest.mark.cpu
-def test_abort_pops_request_states_only_after_ack():
+def test_abort_keeps_request_states_until_generate_cleanup():
+    """Frontend abort must not drop request_states; generate() owns cleanup."""
+
     async def run():
         release = asyncio.Event()
         seen_during_wait: list[int] = []
@@ -215,7 +216,7 @@ def test_abort_pops_request_states_only_after_ack():
 
         release.set()
         await task
-        assert "req-1-aaaa" not in omni.request_states
+        assert "req-1-aaaa" in omni.request_states
 
     asyncio.run(run())
 
@@ -241,6 +242,93 @@ def test_abort_propagates_engine_errors_without_popping_state():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_abort_enqueues_prefix_tokens_from_engine():
+    """AR abort must deliver the generated prefix on the request queue (CPU)."""
+
+    async def run():
+        from vllm.outputs import CompletionOutput
+
+        from vllm_omni.engine.messages import OutputMessage
+
+        queue: asyncio.Queue = asyncio.Queue()
+
+        async def abort_with_prefix(request_ids):
+            rid = request_ids[0]
+            engine_output = OmniRequestOutput(
+                request_id=rid,
+                finished=True,
+                stage_id=0,
+                final_output_type="text",
+                outputs=[
+                    CompletionOutput(
+                        index=0,
+                        text="partial",
+                        token_ids=[7, 8, 9],
+                        cumulative_logprob=None,
+                        logprobs=None,
+                        finish_reason="abort",
+                        stop_reason=None,
+                    )
+                ],
+            )
+            return [
+                OutputMessage(
+                    request_id=rid,
+                    stage_id=0,
+                    replica_id=None,
+                    engine_outputs=engine_output,
+                    metrics=None,
+                    finished=True,
+                    stage_submit_ts=None,
+                )
+            ]
+
+        omni = get_async_omni_instance(fake_abort_request=abort_with_prefix)
+        omni.request_states["req-1-cccc"] = SimpleNamespace(
+            request_id="req-1-cccc",
+            external_request_id="req-1",
+            input_stream_task=None,
+            queue=queue,
+        )
+        await omni._abort(["req-1-cccc"])
+        msg = queue.get_nowait()
+        assert msg.finished is True
+        assert list(msg.engine_outputs.outputs[0].token_ids) == [7, 8, 9]
+        assert msg.engine_outputs.outputs[0].finish_reason == "abort"
+        assert "req-1-cccc" in omni.request_states
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_enqueues_synthetic_finished_when_engine_returns_empty():
+    """Empty abort_async must still unblock generate() with finish_reason=abort."""
+
+    async def run():
+        queue: asyncio.Queue = asyncio.Queue()
+
+        async def empty_abort_async(request_ids):
+            del request_ids
+            return []
+
+        omni = get_async_omni_instance(fake_abort_request=empty_abort_async)
+        omni.request_states["req-1-dddd"] = SimpleNamespace(
+            request_id="req-1-dddd",
+            external_request_id="req-1",
+            input_stream_task=None,
+            queue=queue,
+        )
+        await omni._abort(["req-1-dddd"])
+        msg = queue.get_nowait()
+        assert msg.finished is True
+        assert msg.request_id == "req-1-dddd"
+        assert msg.engine_outputs.outputs[0].finish_reason == "abort"
+        assert "req-1-dddd" 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

@@ -18,11 +18,14 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
     omni = object.__new__(AsyncOmni)
     omni._pause_cond = asyncio.Condition()
     omni._paused = False
+    omni._admitting = 0
     omni._hold_admission_until_resume = False
     omni._sleeping_tags = set()
+    omni._stage_sleeping_tags = {}
     omni._level2_sleeping = False
     omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock())
     omni._final_output_handler = lambda: None
+    omni.reset_mm_cache = AsyncMock()
 
     stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
     omni.engine = SimpleNamespace(
@@ -171,6 +174,33 @@ def test_sleep_blocks_admission_before_engine_core_rpc():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_sleep_waits_for_in_flight_generate_admission():
+    """Sleep must not offload while generate() is still in add_request."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        omni._admitting = 1
+        rpc_started = asyncio.Event()
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                rpc_started.set()
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+        sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+        await asyncio.sleep(0.05)
+        assert not sleep_task.done()
+        assert not rpc_started.is_set()
+        await omni._release_generate_admission()
+        await asyncio.wait_for(sleep_task, timeout=1)
+        assert rpc_started.is_set()
+        assert omni._admitting == 0
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
     async def run() -> None:
@@ -260,6 +290,80 @@ def test_sleep_diffusion_only_skips_engine_core_collective_rpc():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_streaming_wait_for_first_chunk_does_not_block_sleep():
+    """Waiting for the next client chunk must not hold an admission slot."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        rpc_started = asyncio.Event()
+        first_chunk = asyncio.Event()
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                rpc_started.set()
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+        async def wait_then_submit() -> None:
+            await first_chunk.wait()
+            await omni._submit_with_admission(asyncio.sleep(0))
+
+        wait_task = asyncio.create_task(wait_then_submit())
+        sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+        await asyncio.sleep(0.05)
+        assert not wait_task.done()
+        assert omni._admitting == 0
+        await asyncio.wait_for(sleep_task, timeout=1)
+        assert rpc_started.is_set()
+
+        wait_task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await wait_task
+        assert omni._admitting == 0
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_streaming_in_flight_add_blocks_sleep():
+    """Sleep must wait while a streaming ADD/update holds an admission slot."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        add_started = asyncio.Event()
+        add_release = asyncio.Event()
+        rpc_started = asyncio.Event()
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                rpc_started.set()
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+        async def slow_add() -> None:
+            add_started.set()
+            await add_release.wait()
+
+        add_task = asyncio.create_task(omni._submit_with_admission(slow_add()))
+        await add_started.wait()
+        sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+        await asyncio.sleep(0.05)
+        assert not sleep_task.done()
+        assert not rpc_started.is_set()
+        assert omni._admitting == 1
+
+        add_release.set()
+        await asyncio.wait_for(add_task, timeout=1)
+        await asyncio.wait_for(sleep_task, timeout=1)
+        assert rpc_started.is_set()
+        assert omni._admitting == 0
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_wake_up_restores_admission_for_diffusion_only():
     """Pure diffusion must keep sleep → wake → generate (no resume)."""
@@ -308,3 +412,30 @@ def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion():
         assert omni._paused is False
 
     asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_partial_wake_does_not_skip_remaining_sleeping_stage():
+    """sleep(stage_ids=[0]) then wake([0]) must not skip a later wake([1])."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "llm"])
+        await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
+        assert omni._stage_sleeping_tags.keys() == {0, 1}
+
+        await omni.wake_up(stage_ids=[0])
+        assert 0 not in omni._stage_sleeping_tags
+        assert 1 in omni._stage_sleeping_tags
+        assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+
+        omni.collective_rpc.reset_mock()
+        acks = await omni.wake_up(stage_ids=[1])
+        assert acks
+        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"] == [1]
+        assert not omni._sleeping_tags
+        assert not omni._stage_sleeping_tags
+
+    asyncio.run(run())

tests/entrypoints/test_omni_sleep_mode.py

@@ -109,7 +109,7 @@ def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str:
 
 
 async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
-    """Best-effort full wake so the next shared-engine test starts active."""
+    """Best-effort full wake + resume so the next shared-engine test starts active."""
     try:
         await engine.wake_up(stage_ids=stage_ids)
     except Exception as e:
@@ -532,6 +532,7 @@ async def test_pure_diffusion_sleep_wake():
     try:
         await engine.sleep(level=1)
         await engine.wake_up()
+        await engine.resume_generation()
         async for _ in engine.generate(
             "test",
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),

tests/entrypoints/test_stage_utils.py

@@ -216,6 +216,20 @@ def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytes
     assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4"
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_resolve_stage_physical_devices_captured_unset_ignores_live_env(monkeypatch: pytest.MonkeyPatch):
+    """Parallel init captures an unset CUDA_VISIBLE_DEVICES as None.
+
+    A sibling thread may then narrow the live env for spawn. The captured
+    None must not fall back to that live value, or logical device 0 would
+    map onto the other stage's GPU.
+    """
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
+    assert resolve_stage_physical_devices(0, "0", visible_baseline=None) == "0"
+    assert resolve_stage_physical_devices(1, "1", visible_baseline=None) == "1"
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_raises_on_non_numeric():

tests/worker/test_worker_base.py

@@ -199,7 +199,14 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
     worker = object.__new__(OmniGPUWorkerBase)
     worker.rank = 0
     worker.device = "cuda:0"
-    monkeypatch.setattr(base, "current_omni_platform", _fake_platform())
+    order: list[str] = []
+    platform = SimpleNamespace(
+        get_current_memory_usage=lambda device: 0,
+        empty_cache=lambda: order.append("empty_cache"),
+        synchronize=lambda: order.append("sync"),
+        collect=lambda: None,
+    )
+    monkeypatch.setattr(base, "current_omni_platform", platform)
 
     calls = {}
 
@@ -209,12 +216,17 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
             return cls()
 
         def sleep(self, offload_tags):
+            order.append("allocator_sleep")
             calls["offload_tags"] = offload_tags
 
     monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator)
 
     assert OmniGPUWorkerBase.sleep(worker, level=level) is True
     assert calls["offload_tags"] == expected_tags
+    assert order[0] == "sync"
+    assert "allocator_sleep" in order
+    assert order.index("sync") < order.index("allocator_sleep")
+    assert "empty_cache" not in order
 
 
 def test_wake_up_forwards_tags_to_allocator(monkeypatch):

vllm_omni/engine/async_omni_engine.py

@@ -67,6 +67,7 @@ from vllm_omni.engine.messages import (
     EngineQueueMessage,
     ErrorMessage,
     InteractionMessage,
+    OutputMessage,
     StageSubmissionMessage,
 )
 from vllm_omni.engine.orchestrator import Orchestrator
@@ -1920,12 +1921,16 @@ class AsyncOmniEngine:
         self,
         request_ids: list[str],
         timeout: float | None = None,
-    ) -> None:
+    ) -> list[OutputMessage]:
         """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.
+
+        Returns:
+            Final-stage AR abort ``OutputMessage`` list carrying partial
+            tokens generated before abort (empty for diffusion / no OP state).
         """
         if not request_ids or getattr(self, "_shutdown_called", False):
             return
@@ -1959,6 +1964,7 @@ class AsyncOmniEngine:
             raise
         if not result_msg.success:
             raise RuntimeError(result_msg.error or "abort failed")
+        return list(result_msg.abort_outputs or [])
 
     def submit_interaction(
         self,

vllm_omni/engine/messages.py

@@ -48,17 +48,6 @@ class AbortRequestMessage(EngineQueueMessage, kw_only=True):
     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):
     type: Literal["interaction"] = "interaction"
     request_id: str
@@ -113,6 +102,20 @@ class OutputMessage(EngineQueueMessage, kw_only=True):
     stage_submit_ts: float | None = None
 
 
+class AbortResultMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["abort_result"] = "abort_result"
+    rpc_id: str
+    success: bool
+    error: str | None = None
+    # Final-stage AR abort outputs (partial tokens) for frontend delivery.
+    # Empty for diffusion / requests with no output-processor state.
+    abort_outputs: list[OutputMessage] | None = None
+
+    @property
+    def rpc_correlation_key(self) -> tuple[str, str]:
+        return ("abort", self.rpc_id)
+
+
 class StageMetricsMessage(EngineQueueMessage, kw_only=True):
     type: Literal["stage_metrics"] = "stage_metrics"
     request_id: str

vllm_omni/engine/orchestrator.py

@@ -871,14 +871,19 @@ class Orchestrator:
         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.
+
+        AR final-stage abort outputs (partial tokens) are attached to the
+        result message when ``rpc_id`` is set, or enqueued on
+        ``output_async_queue`` for fire-and-forget aborts.
         """
         request_ids = msg.request_ids
         error: str | None = None
+        abort_outputs: list[OutputMessage] = []
         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)
+            abort_outputs = 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)
@@ -893,8 +898,12 @@ class Orchestrator:
                     rpc_id=msg.rpc_id,
                     success=error is None,
                     error=error,
+                    abort_outputs=abort_outputs or None,
                 )
             )
+        elif abort_outputs:
+            for output_msg in abort_outputs:
+                await self.output_async_queue.put(output_msg)
 
     async def _handle_interaction(self, msg: InteractionMessage) -> None:
         """Handle a midway interaction for an active streaming diffusion request."""
@@ -934,13 +943,53 @@ class Orchestrator:
                 )
             )
 
-    async def _abort_request_ids(self, request_ids: list[str]) -> None:
-        """Forward abort requests to all stage pools."""
+    async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]:
+        """Forward abort requests to all stage pools.
+
+        Collects final-stage AR abort outputs (partial tokens) while request
+        state is still available for stage-id filtering. Diffusion pools do
+        not contribute abort outputs.
+        """
         if not request_ids:
-            return
+            return []
+        abort_outputs: list[OutputMessage] = []
         for pool in self.stage_pools:
-            await pool.abort_requests(request_ids)
+            stage_outputs = await pool.abort_requests(request_ids) or []
+            if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion":
+                final_output_type = getattr(pool.stage_client, "final_output_type", None) or "text"
+                for orch_req_id, request_output in stage_outputs:
+                    req_state = self.request_states.get(orch_req_id)
+                    if req_state is None:
+                        continue
+                    final_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
+                    if pool.stage_id not in final_stage_ids:
+                        continue
+                    engine_output = OmniRequestOutput.from_stage_output(
+                        request_output,
+                        request_id=orch_req_id,
+                        finished=True,
+                        stage_id=pool.stage_id,
+                        replica_id=pool.get_bound_replica_id(orch_req_id),
+                        final_output_type=final_output_type,
+                    )
+                    abort_outputs.append(
+                        OutputMessage(
+                            request_id=orch_req_id,
+                            stage_id=pool.stage_id,
+                            replica_id=pool.get_bound_replica_id(orch_req_id),
+                            engine_outputs=engine_output,
+                            metrics=None,
+                            finished=False,
+                            stage_submit_ts=req_state.stage_submit_ts.get(pool.stage_id),
+                        )
+                    )
             pool.release_bindings(request_ids)
+        last_index_by_req: dict[str, int] = {}
+        for index, output_msg in enumerate(abort_outputs):
+            last_index_by_req[output_msg.request_id] = index
+        for index, output_msg in enumerate(abort_outputs):
+            output_msg.finished = index == last_index_by_req[output_msg.request_id]
+        return abort_outputs
 
     def _release_request_bindings(self, request_ids: list[str]) -> None:
         """Release all stage-local route bindings for the given request ids."""
@@ -1688,7 +1737,7 @@ class Orchestrator:
         *,
         abort: bool = False,
         close_duplex_sessions: bool = False,
-    ) -> None:
+    ) -> list[OutputMessage]:
         """Release pool bindings and logical request state for the given ids.
 
         CFG-aware: cleaning a parent releases its tracker state and pulls its
@@ -1697,9 +1746,13 @@ class Orchestrator:
         can never complete). Every teardown path (stage error, abort, replica
         loss, membership unregister) funnels through here, so tracker state
         cannot outlive its requests.
+
+        Returns:
+            Final-stage AR abort ``OutputMessage`` list when ``abort=True``;
+            otherwise an empty list.
         """
         if not request_ids:
-            return
+            return []
 
         cleanup_ids = list(dict.fromkeys(request_ids))
         batch = set(cleanup_ids)
@@ -1741,9 +1794,10 @@ class Orchestrator:
                 cleanup_ids.extend(stale_request_ids)
             cleanup_ids = list(dict.fromkeys(cleanup_ids))
 
+        abort_outputs: list[OutputMessage] = []
         try:
             if abort:
-                await self._abort_request_ids(cleanup_ids)
+                abort_outputs = await self._abort_request_ids(cleanup_ids)
             self._release_request_bindings(cleanup_ids)
             for request_id in cleanup_ids:
                 self._pd_kv_params.pop(request_id, None)
@@ -1759,6 +1813,7 @@ class Orchestrator:
             raise
         if closing_session_ids and self.duplex_control_plane is not None:
             self.duplex_control_plane.finalize_closed_sessions(closing_session_ids)
+        return abort_outputs
 
     async def _apply_raw_terminal_stage_finish(
         self,

vllm_omni/engine/stage_pool.py

@@ -1189,14 +1189,19 @@ class StagePool:
 
     # ---- Stage-local control plane ----
 
-    async def abort_requests(self, request_ids: list[str]) -> None:
+    async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]:
         """Abort the given requests in this stage pool.
 
         Request-bound abort routing stays inside the pool because route affinity
         (``request_id -> replica_id``) is pool-owned.
+
+        For AR stages, collect terminal abort ``RequestOutput`` objects from
+        output-processor state *before* EngineCore abort so partial tokens
+        generated so far are preserved. Diffusion stages return no abort
+        outputs (whole-sample retry remains the recovery path).
         """
         if not request_ids:
-            return
+            return []
 
         request_ids_by_replica: dict[int, list[str]] = {}
         for request_id in request_ids:
@@ -1206,18 +1211,45 @@ class StagePool:
                 continue
             request_ids_by_replica.setdefault(replica_id, []).append(request_id)
 
+        abort_outputs: list[tuple[str, Any]] = []
+        is_diffusion = self.stage_type == "diffusion"
         for replica_id, replica_request_ids in request_ids_by_replica.items():
+            # Orchestrator ids are OP external ids; EngineCore may use a
+            # different internal request_id assigned by InputProcessor.
+            # Always key collected outputs by the orchestrator id we passed
+            # (internal=False), never by RequestOutput.request_id — some stages
+            # still surface EngineCore/internal ids and would drop the prefix.
+            engine_abort_ids = list(replica_request_ids)
+            if not is_diffusion and self._output_processor is not None:
+                collect = getattr(self._output_processor, "abort_requests_collecting_outputs", None)
+                if collect is not None:
+                    engine_abort_ids = []
+                    for orch_req_id in replica_request_ids:
+                        collected_ids, stage_outputs = collect(
+                            [orch_req_id],
+                            internal=False,
+                            commit_state=False,
+                        )
+                        for req_out in stage_outputs:
+                            abort_outputs.append((orch_req_id, req_out))
+                        if collected_ids:
+                            engine_abort_ids.extend(collected_ids)
+                        else:
+                            engine_abort_ids.append(orch_req_id)
+                else:
+                    aborted = self._output_processor.abort_requests(replica_request_ids, internal=False)
+                    if aborted:
+                        engine_abort_ids = list(aborted)
             client = self.clients[replica_id]
             if client is None:
                 continue
-            await client.abort_requests_async(replica_request_ids)
+            await client.abort_requests_async(engine_abort_ids)
+            if not is_diffusion and self._output_processor is not None:
+                commit = getattr(self._output_processor, "commit_aborted_request_state", None)
+                if callable(commit):
+                    commit(replica_request_ids, internal=False)
 
-        # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that
-        # would otherwise leak — aborted requests never produce a final
-        # EngineCoreOutput, so process_outputs() never fires its cleanup path.
-        all_aborted = [rid for ids in request_ids_by_replica.values() for rid in ids]
-        if all_aborted and self._output_processor is not None:
-            self._output_processor.abort_requests(all_aborted, internal=True)
+        return abort_outputs
 
     async def collective_rpc(
         self,
@@ -1227,7 +1259,12 @@ class StagePool:
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any] | Any:
-        """Dispatch a stage-scoped control-plane RPC to one physical route."""
+        """Dispatch a stage-scoped control-plane RPC to one physical route.
+
+        EngineCore control methods (``pause_scheduler``, ``resume_scheduler``,
+        ``sleep``, ``wake_up``) propagate worker exceptions. Capability probes
+        and other RPCs still return ``{"supported": False, "error": ...}``.
+        """
         args = tuple(args or ())
         kwargs = dict(kwargs or {})
         client = self.clients[replica_id]
@@ -1258,9 +1295,15 @@ class StagePool:
                 replica_id,
                 method,
             )
+            if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS:
+                raise
+            if isinstance(exc, TimeoutError):
+                error = f"{type(exc).__name__}: {method} timed out after {timeout}s"
+            else:
+                error = str(exc) or repr(exc)
             return {
                 "supported": False,
-                "error": str(exc),
+                "error": error,
             }
 
     def shutdown_replica(self, replica_id: int) -> None:

vllm_omni/entrypoints/async_omni.py

@@ -17,7 +17,7 @@ from vllm import TokensPrompt
 from vllm.engine.protocol import EngineClient, StreamingInput
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
-from vllm.outputs import PoolingRequestOutput
+from vllm.outputs import CompletionOutput, PoolingRequestOutput
 from vllm.plugins.io_processors import get_io_processor
 from vllm.pooling_params import PoolingParams
 from vllm.renderers.inputs.preprocess import extract_prompt_components
@@ -150,12 +150,19 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
+        # In-flight EngineCore submits (non-streaming add_request, or each
+        # streaming ADD/update/final marker). sleep() waits for this to hit
+        # zero so a pipelined request cannot race into EngineCore during
+        # drain/offload. Streaming generate does not hold a slot while
+        # waiting for the next client chunk.
+        self._admitting: int = 0
         # 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._stage_sleeping_tags: dict[int, set[str]] = {}
         self._level2_sleeping: bool = False
         self._duplex_request_client: DuplexRequestClient | None = None
         self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue()
@@ -513,31 +520,40 @@ class AsyncOmni(EngineClient, OmniBase):
         external_request_id = request_id
         request_id = self._get_unique_request_id(external_request_id)
 
-        # Wait until generation is resumed if the engine is paused
+        # Wait until generation is resumed if the engine is paused. Non-streaming
+        # generate holds an admission slot until add_request completes so sleep()
+        # cannot race a just-unblocked generate into EngineCore during offload.
+        # Streaming generate does **not** hold the slot while waiting for the
+        # first client chunk; the input pump acquires it immediately before each
+        # EngineCore ADD/update.
+        streaming_input = isinstance(prompt, AsyncGenerator)
         async with self._pause_cond:
             await self._pause_cond.wait_for(lambda: not self._paused)
+            if not streaming_input:
+                self._admitting = getattr(self, "_admitting", 0) + 1
+        admitting = not streaming_input
 
         logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")
 
-        _sleeping_tags = getattr(self, "_sleeping_tags", None)
-        if _sleeping_tags:
-            raise RuntimeError(
-                f"Generation rejected: Engine is partially or fully asleep. "
-                f"Currently sleeping tags: {list(_sleeping_tags)}. "
-                f"Please perform a full wake_up before generating."
-            )
-
-        # Reject diffusion list-prompt early with a clear API error.
-        if isinstance(prompt, list) and any(
-            getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
-        ):
-            raise ValueError(
-                "Diffusion stages accept only a single prompt per request. "
-                "Submit multiple independent requests to use scheduler batching."
-            )
-
         input_stream_task: asyncio.Task | None = None
         try:
+            _sleeping_tags = getattr(self, "_sleeping_tags", None)
+            if _sleeping_tags:
+                raise RuntimeError(
+                    f"Generation rejected: Engine is partially or fully asleep. "
+                    f"Currently sleeping tags: {list(_sleeping_tags)}. "
+                    f"Please perform a full wake_up before generating."
+                )
+
+            # Reject diffusion list-prompt early with a clear API error.
+            if isinstance(prompt, list) and any(
+                getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
+            ):
+                raise ValueError(
+                    "Diffusion stages accept only a single prompt per request. "
+                    "Submit multiple independent requests to use scheduler batching."
+                )
+
             # Start final output dispatcher on the first call to generate()
             self._final_output_handler()
 
@@ -598,8 +614,11 @@ class AsyncOmni(EngineClient, OmniBase):
                 req_sp_list[p_id] = self._prepare_prefill_sampling_params(request_id, req_sp_list[p_id])
 
             # Add request(s) to stage 0. For streaming inputs, submit
-            # chunks incrementally through streaming_update.
-            if isinstance(prompt, AsyncGenerator):
+            # chunks incrementally through streaming_update. The helper
+            # returns as soon as the pump task is created; each ADD/update
+            # takes its own admission slot inside the pump.
+            if streaming_input:
+                first_chunk_submitted = asyncio.get_running_loop().create_future()
                 input_stream_task = await self._add_streaming_input_request(
                     request_id=request_id,
                     input_stream=prompt,
@@ -608,7 +627,9 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_output_stage_ids=final_output_stage_ids,
                     arrival_time=wall_start_ts,
                     lora_request=lora_request,
+                    first_chunk_submitted=first_chunk_submitted,
                 )
+                await first_chunk_submitted
             else:
                 await self.engine.add_request_async(
                     request_id=request_id,
@@ -622,6 +643,9 @@ class AsyncOmni(EngineClient, OmniBase):
             submit_ts = time.time()
             req_state.metrics.stage_first_ts[0] = submit_ts
             req_start_ts[request_id] = submit_ts
+            if admitting:
+                await self._release_generate_admission()
+                admitting = False
             # Refresh gauges on arrival.
             self._publish_request_gauges(len(self.request_states))
 
@@ -640,18 +664,7 @@ class AsyncOmni(EngineClient, OmniBase):
 
             logger.debug(f"[AsyncOmni] Request {request_id} completed")
 
-            # The input pump can outlive a normally-completed generation
-            # (e.g. a realtime client that keeps streaming after the final
-            # output); a live pump keeps issuing streaming updates for a
-            # request the orchestrator no longer tracks.
-            if input_stream_task is not None and not input_stream_task.done():
-                input_stream_task.cancel()
-
-            self._log_summary_and_cleanup(request_id)
-
         except (asyncio.CancelledError, GeneratorExit):
-            if input_stream_task is not None and not input_stream_task.done():
-                input_stream_task.cancel()
             self._record_request_failure_once(request_id, reason="client_disconnect")
             await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} aborted.")
@@ -661,6 +674,28 @@ class AsyncOmni(EngineClient, OmniBase):
             await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}")
             raise
+        finally:
+            if input_stream_task is not None and not input_stream_task.done():
+                input_stream_task.cancel()
+            if admitting:
+                await self._release_generate_admission()
+            self._log_summary_and_cleanup(request_id)
+
+    async def _release_generate_admission(self) -> None:
+        """Drop one in-flight generate admission slot held across add_request."""
+        async with self._pause_cond:
+            self._admitting = max(getattr(self, "_admitting", 1) - 1, 0)
+            self._pause_cond.notify_all()
+
+    async def _submit_with_admission(self, awaitable):
+        """Wait for resume, hold one admission slot for a single EngineCore submit."""
+        async with self._pause_cond:
+            await self._pause_cond.wait_for(lambda: not self._paused)
+            self._admitting = getattr(self, "_admitting", 0) + 1
+        try:
+            return await awaitable
+        finally:
+            await self._release_generate_admission()
 
     async def _add_streaming_input_request(
         self,
@@ -672,6 +707,7 @@ class AsyncOmni(EngineClient, OmniBase):
         final_output_stage_ids: Sequence[int],
         arrival_time: float,
         lora_request: Any = None,
+        first_chunk_submitted: asyncio.Future[None] | None = None,
     ) -> asyncio.Task:
         """Submit a streaming input generator as incremental stage-0 updates."""
         if not sampling_params_list:
@@ -688,6 +724,10 @@ class AsyncOmni(EngineClient, OmniBase):
             stage0_params = stage0_params.clone()
             stage0_params.skip_clone = True
 
+        def _mark_first_chunk_submitted() -> None:
+            if first_chunk_submitted is not None and not first_chunk_submitted.done():
+                first_chunk_submitted.set_result(None)
+
         async def handle_inputs() -> None:
             nonlocal has_submitted_first_chunk
             cancelled = False
@@ -701,29 +741,34 @@ class AsyncOmni(EngineClient, OmniBase):
                     prompt_text, _, _ = extract_prompt_components(self.model_config, chunk_prompt)
 
                     if not has_submitted_first_chunk:
-                        await self.engine.add_request_async(
-                            request_id=request_id,
-                            prompt=chunk_prompt,
-                            prompt_text=prompt_text,
-                            sampling_params_list=chunk_sampling_params_list,
-                            final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
-                            lora_request=lora_request,
-                            resumable=True,
+                        await self._submit_with_admission(
+                            self.engine.add_request_async(
+                                request_id=request_id,
+                                prompt=chunk_prompt,
+                                prompt_text=prompt_text,
+                                sampling_params_list=chunk_sampling_params_list,
+                                final_stage_id=final_stage_id,
+                                final_output_stage_ids=final_output_stage_ids,
+                                arrival_time=arrival_time,
+                                lora_request=lora_request,
+                                resumable=True,
+                            )
                         )
                         has_submitted_first_chunk = True
+                        _mark_first_chunk_submitted()
                     else:
-                        await self.engine.add_streaming_update_async(
-                            request_id=request_id,
-                            prompt=chunk_prompt,
-                            prompt_text=prompt_text,
-                            sampling_params_list=chunk_sampling_params_list,
-                            final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
-                            lora_request=lora_request,
-                            resumable=True,
+                        await self._submit_with_admission(
+                            self.engine.add_streaming_update_async(
+                                request_id=request_id,
+                                prompt=chunk_prompt,
+                                prompt_text=prompt_text,
+                                sampling_params_list=chunk_sampling_params_list,
+                                final_stage_id=final_stage_id,
+                                final_output_stage_ids=final_output_stage_ids,
+                                arrival_time=arrival_time,
+                                lora_request=lora_request,
+                                resumable=True,
+                            )
                         )
             except (asyncio.CancelledError, GeneratorExit):
                 cancelled = True
@@ -738,37 +783,47 @@ class AsyncOmni(EngineClient, OmniBase):
                     )
                 )
             finally:
-                if not cancelled:
-                    # Send empty final request to indicate that inputs have
-                    # finished. Don't send if canceled (session was aborted).
-                    final_sampling_params_list = list(sampling_params_list)
-                    final_sampling_params_list[0] = stage0_params
-                    final_prompt = TokensPrompt(prompt_token_ids=[0])
+                try:
+                    if not cancelled:
+                        # Send empty final request to indicate that inputs have
+                        # finished. Don't send if canceled (session was aborted).
+                        final_sampling_params_list = list(sampling_params_list)
+                        final_sampling_params_list[0] = stage0_params
+                        final_prompt = TokensPrompt(prompt_token_ids=[0])
 
-                    if has_submitted_first_chunk:
-                        await self.engine.add_streaming_update_async(
-                            request_id=request_id,
-                            prompt=final_prompt,
-                            prompt_text=None,
-                            sampling_params_list=final_sampling_params_list,
-                            final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
-                            lora_request=lora_request,
-                            resumable=False,
-                        )
-                    else:
-                        await self.engine.add_request_async(
-                            request_id=request_id,
-                            prompt=final_prompt,
-                            prompt_text=None,
-                            sampling_params_list=final_sampling_params_list,
-                            final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
-                            lora_request=lora_request,
-                            resumable=False,
-                        )
+                        if has_submitted_first_chunk:
+                            await self._submit_with_admission(
+                                self.engine.add_streaming_update_async(
+                                    request_id=request_id,
+                                    prompt=final_prompt,
+                                    prompt_text=None,
+                                    sampling_params_list=final_sampling_params_list,
+                                    final_stage_id=final_stage_id,
+                                    final_output_stage_ids=final_output_stage_ids,
+                                    arrival_time=arrival_time,
+                                    lora_request=lora_request,
+                                    resumable=False,
+                                )
+                            )
+                        else:
+                            await self._submit_with_admission(
+                                self.engine.add_request_async(
+                                    request_id=request_id,
+                                    prompt=final_prompt,
+                                    prompt_text=None,
+                                    sampling_params_list=final_sampling_params_list,
+                                    final_stage_id=final_stage_id,
+                                    final_output_stage_ids=final_output_stage_ids,
+                                    arrival_time=arrival_time,
+                                    lora_request=lora_request,
+                                    resumable=False,
+                                )
+                            )
+                            has_submitted_first_chunk = True
+                finally:
+                    # Unblock generate() even on cancel / empty stream / submit
+                    # failure so it can observe a terminal abort or empty result.
+                    _mark_first_chunk_submitted()
 
         input_stream_task = asyncio.create_task(handle_inputs())
         req_state.input_stream_task = input_stream_task
@@ -1129,22 +1184,76 @@ class AsyncOmni(EngineClient, OmniBase):
         await self._abort(internal_req_ids)
 
     async def _abort(self, request_ids: list[str]) -> None:
-        """Abort request IDs via the engine and clean frontend state after ack.
+        """Abort request IDs via the engine and enqueue terminal abort outputs.
 
-        Waits for orchestrator abort acknowledgment before popping
-        ``request_states`` so generate() cleanup stays consistent with
-        backend binding/request teardown. Orchestrator abort errors propagate.
+        Waits for orchestrator abort acknowledgment, enqueues any AR terminal
+        abort outputs (partial tokens) into each request's asyncio queue, then
+        cancels the input pump. Frontend ``request_states`` stay registered so
+        ``generate()`` can consume the terminal message in
+        ``_process_orchestrator_results`` and run normal cleanup.
+
+        When ``abort_async`` returns no output for an active request (OP not
+        registered yet, unbound replica, or orchestrator id drop), enqueue a
+        synthetic finished abort so ``generate()`` cannot hang on ``queue.get``.
         """
-        await self.engine.abort_async(request_ids)
+        abort_outputs = await self.engine.abort_async(request_ids) or []
+        delivered: set[str] = set()
+        for output_msg in abort_outputs:
+            req_id = getattr(output_msg, "request_id", None)
+            if req_id is None:
+                continue
+            state = self.request_states.get(req_id)
+            if state is None:
+                logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id)
+                continue
+            await state.queue.put(output_msg)
+            delivered.add(req_id)
+        for rid in request_ids:
+            state = self.request_states.get(rid)
+            if state is not None and rid not in delivered:
+                queue = getattr(state, "queue", None)
+                if queue is not None:
+                    await state.queue.put(self._synthetic_abort_output_message(rid))
+                    delivered.add(rid)
         for rid in request_ids:
             self._record_request_failure_once(rid, reason="client_abort")
-            state = self.request_states.pop(rid, None)
+            state = self.request_states.get(rid)
             input_stream_task = getattr(state, "input_stream_task", None)
             if input_stream_task is not None and not input_stream_task.done():
                 input_stream_task.cancel()
         if self.log_stats:
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 
+    @staticmethod
+    def _synthetic_abort_output_message(request_id: str) -> OutputMessage:
+        """Terminal abort OutputMessage used when the engine returned none."""
+        engine_output = OmniRequestOutput(
+            request_id=request_id,
+            finished=True,
+            stage_id=0,
+            final_output_type="text",
+            outputs=[
+                CompletionOutput(
+                    index=0,
+                    text="",
+                    token_ids=[],
+                    cumulative_logprob=None,
+                    logprobs=None,
+                    finish_reason="abort",
+                    stop_reason=None,
+                )
+            ],
+        )
+        return OutputMessage(
+            request_id=request_id,
+            stage_id=0,
+            replica_id=None,
+            engine_outputs=engine_output,
+            metrics=None,
+            finished=True,
+            stage_submit_ts=None,
+        )
+
     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)
@@ -1168,6 +1277,47 @@ class AsyncOmni(EngineClient, OmniBase):
                 ar_stage_ids.append(sid)
         return ar_stage_ids, diffusion_stage_ids
 
+    def _sleeping_tags_for_stages(self, stage_ids: list[int]) -> set[str]:
+        """Union of sleeping tags recorded for the given stages."""
+        per_stage = getattr(self, "_stage_sleeping_tags", None) or {}
+        tags: set[str] = set()
+        for sid in stage_ids:
+            tags.update(per_stage.get(sid, set()))
+        return tags
+
+    def _refresh_union_sleeping_tags(self) -> None:
+        """Keep ``_sleeping_tags`` as the engine-wide union of per-stage tags."""
+        per_stage = getattr(self, "_stage_sleeping_tags", None) or {}
+        union: set[str] = set()
+        for stage_tags in per_stage.values():
+            union.update(stage_tags)
+        self._sleeping_tags = union
+
+    def _record_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None:
+        per_stage = getattr(self, "_stage_sleeping_tags", None)
+        if per_stage is None:
+            per_stage = {}
+            self._stage_sleeping_tags = per_stage
+        tag_set = set(tags)
+        for sid in stage_ids:
+            per_stage.setdefault(sid, set()).update(tag_set)
+        self._refresh_union_sleeping_tags()
+
+    def _clear_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None:
+        per_stage = getattr(self, "_stage_sleeping_tags", None)
+        if per_stage is None:
+            self._sleeping_tags = set()
+            return
+        tag_set = set(tags)
+        for sid in stage_ids:
+            remaining = per_stage.get(sid)
+            if remaining is None:
+                continue
+            remaining.difference_update(tag_set)
+            if not remaining:
+                per_stage.pop(sid, None)
+        self._refresh_union_sleeping_tags()
+
     async def pause_generation(
         self,
         *,
@@ -1266,11 +1416,24 @@ class AsyncOmni(EngineClient, OmniBase):
         return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages)
 
     async def reset_mm_cache(self) -> None:
-        """Reset the multi-modal cache for all stages.
+        """Reset the frontend (P0) multimodal processor cache.
 
-        TODO: Forward to Orchestrator process via message.
+        ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache.
+        Clearing P0 avoids hash-only follow-up requests after that reset.
         """
-        logger.warning("[AsyncOmni] reset_mm_cache not yet supported with Orchestrator process")
+        processor = getattr(self, "input_processor", None)
+        if processor is None:
+            processor = getattr(self.engine, "input_processor", None)
+        cache = getattr(processor, "mm_processor_cache", None)
+        if cache is None:
+            logger.debug("[AsyncOmni] reset_mm_cache: no frontend mm_processor_cache")
+            return
+        for name in ("clear", "reset", "clear_cache"):
+            fn = getattr(cache, name, None)
+            if callable(fn):
+                fn()
+                return
+        logger.debug("[AsyncOmni] reset_mm_cache: cache has no clear/reset method")
 
     async def reset_encoder_cache(self) -> None:
         """Reset the encoder cache for all stages.
@@ -1314,10 +1477,15 @@ class AsyncOmni(EngineClient, OmniBase):
         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.
+        # _pause_cond during the drain/offload window. Wait until generate()
+        # coroutines that already passed the pause check have submitted (or
+        # failed) so EngineCore does not see ADD frames while sleeping.
         async with self._pause_cond:
             self._paused = True
+            await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0)
+
+        # P0 sender cache must drop hashes before EngineCore.sleep clears P1.
+        await self.reset_mm_cache()
 
         self._final_output_handler()
         ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
@@ -1353,9 +1521,10 @@ class AsyncOmni(EngineClient, OmniBase):
         if diffusion_stage_ids:
             final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level))
 
-        if not hasattr(self, "_sleeping_tags"):
-            self._sleeping_tags = set()
-        self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value])
+        self._record_stage_sleep(
+            ar_stage_ids + diffusion_stage_ids,
+            [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value],
+        )
         if level == 2:
             self._level2_sleeping = True
         return final_acks
@@ -1399,7 +1568,12 @@ class AsyncOmni(EngineClient, OmniBase):
                 "Use sleep(level=1) instead, which offloads weights to CPU RAM "
                 "and supports fast DMA restore."
             )
-        _current_tags = getattr(self, "_sleeping_tags", set())
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        target_stage_ids = ar_stage_ids + diffusion_stage_ids
+        _current_tags = self._sleeping_tags_for_stages(target_stage_ids)
+        per_stage = getattr(self, "_stage_sleeping_tags", None) or {}
+        if not _current_tags and not per_stage:
+            _current_tags = set(getattr(self, "_sleeping_tags", set()))
         if tags is None:
             requested_tags = list(_current_tags)
         else:
@@ -1408,8 +1582,6 @@ class AsyncOmni(EngineClient, OmniBase):
             logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
             return []
 
-        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)
@@ -1433,9 +1605,7 @@ class AsyncOmni(EngineClient, OmniBase):
         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"):
-                self._sleeping_tags.discard(t)
+        self._clear_stage_sleep(target_stage_ids, requested_tags)
         # Only clear the level-2 flag once all tags are warm, in case partial
         # wake support (e.g. tags=["kv_cache"] only) is added in the future.
         if not getattr(self, "_sleeping_tags", None):

vllm_omni/entrypoints/stage_utils.py

@@ -10,19 +10,26 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
+# Default for ``visible_baseline``: read the live device-control env var.
+# Explicit ``None`` means the caller captured an unset env and must not fall
+# back to a later thread's temporary CUDA_VISIBLE_DEVICES.
+_LIVE_DEVICE_ENV = object()
+
 
 def resolve_stage_physical_devices(
     stage_id: int,
     devices: str | int | None,
     *,
-    visible_baseline: str | None = None,
+    visible_baseline: str | None | object = _LIVE_DEVICE_ENV,
 ) -> str | None:
     """Map logical stage devices to physical IDs without mutating process env.
 
-    When ``visible_baseline`` is provided it is used for logical-to-physical
-    mapping instead of the current device-control env var. This keeps parallel
-    stage initialization correct even if another thread has temporarily
-    narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage.
+    ``visible_baseline`` is the device-control env captured at the start of
+    replica init. Pass the captured string, or ``None`` when that env was
+    unset. Omitting the argument reads the live env (single-threaded callers).
+
+    Parallel stage init must pass the captured value: another thread may have
+    already narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage.
     """
     if devices in (None, "cpu"):
         return None
@@ -32,7 +39,7 @@ def resolve_stage_physical_devices(
 
     device_list = _parse_device_list(devices)
     env_var = current_omni_platform.device_control_env_var
-    visible_devices = visible_baseline if visible_baseline is not None else os.environ.get(env_var)
+    visible_devices = os.environ.get(env_var) if visible_baseline is _LIVE_DEVICE_ENV else visible_baseline
     if visible_devices is not None:
         visible_device_list = _parse_device_list(visible_devices)
         device_list = _map_device_list(stage_id, device_list, visible_device_list)
@@ -43,7 +50,7 @@ def set_stage_devices(
     stage_id: int,
     devices: str | int | None,
     *,
-    visible_baseline: str | None = None,
+    visible_baseline: str | None | object = _LIVE_DEVICE_ENV,
 ) -> str | None:
     """Configure per-stage device visibility and current device (CUDA or NPU).
 

vllm_omni/outputs/output_processor.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from collections.abc import Iterable
 from dataclasses import fields as dataclass_fields
 from typing import Any
 
@@ -432,6 +433,116 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         return self._native_text_metrics_by_request.pop(request_id, {})
 
     def abort_requests(self, request_ids, internal: bool) -> list[str]:
+        aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal)
+        return aborted_ids
+
+    def abort_requests_collecting_outputs(
+        self,
+        request_ids: Iterable[str],
+        *,
+        internal: bool,
+        commit_state: bool = True,
+    ) -> tuple[list[str], list[RequestOutput | PoolingRequestOutput]]:
+        """Abort requests and return terminal abort outputs with partial tokens.
+
+        Mirrors upstream ``OutputProcessor.abort_requests``, but also returns
+        abort ``RequestOutput`` objects when ``queue`` is ``None`` (Omni stage
+        pools register OP state without a collector queue).
+
+        Terminal abort outputs use ``RequestOutputKind.CUMULATIVE`` so DELTA
+        streaming requests still surface the full prefix generated so far.
+        ``new_token_ids`` is the detokenizer prefix: Omni's
+        ``make_request_output`` writes that list onto ``CompletionOutput.token_ids``.
+
+        When ``commit_state`` is False, processor mappings stay in place so a
+        failed physical EngineCore abort can retry with the same prefix.
+        """
+        request_ids = list(request_ids)
+        if commit_state:
+            for request_id in request_ids:
+                if internal:
+                    req_state = self.request_states.get(request_id)
+                    if req_state is not None:
+                        self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
+                else:
+                    self._native_text_metrics_by_request.pop(request_id, None)
+
+        internal_req_ids: list[str] = []
+        for request_id in request_ids:
+            if internal:
+                internal_req_ids.append(request_id)
+                if commit_state:
+                    if req_state := self.request_states.get(request_id):
+                        external_req_id = req_state.external_req_id
+                        internal_ids = self.external_req_ids.get(external_req_id)
+                        if internal_ids is not None:
+                            try:
+                                internal_ids.remove(request_id)
+                            except ValueError:
+                                pass
+                            if not internal_ids:
+                                self.external_req_ids.pop(external_req_id, None)
+            elif commit_state:
+                if internal_ids := self.external_req_ids.pop(request_id, []):
+                    internal_req_ids.extend(internal_ids)
+            else:
+                internal_req_ids.extend(self.external_req_ids.get(request_id, []))
+
+        request_ids_to_abort: list[str] = []
+        abort_outputs: list[RequestOutput | PoolingRequestOutput] = []
+        for request_id in internal_req_ids:
+            req_state = (
+                self.request_states.pop(request_id, None) if commit_state else self.request_states.get(request_id)
+            )
+            if req_state is not None:
+                if commit_state:
+                    self.lora_states.request_finished(request_id, req_state.lora_name)
+                request_ids_to_abort.append(request_id)
+                original_kind = req_state.output_kind
+                req_state.output_kind = RequestOutputKind.CUMULATIVE
+                detok = req_state.detokenizer
+                new_token_ids = list(detok.output_token_ids) if detok is not None else []
+                saved_children = None
+                if not commit_state and req_state.parent_req is not None:
+                    saved_children = set(req_state.parent_req.child_requests)
+                try:
+                    request_output = req_state.make_request_output(
+                        new_token_ids=new_token_ids,
+                        pooling_output=None,
+                        finish_reason=FinishReason.ABORT,
+                        stop_reason=None,
+                        kv_transfer_params=None,
+                    )
+                finally:
+                    req_state.output_kind = original_kind
+                    if saved_children is not None and req_state.parent_req is not None:
+                        req_state.parent_req.child_requests = saved_children
+                if request_output is not None:
+                    abort_outputs.append(request_output)
+                    if commit_state and req_state.queue is not None:
+                        req_state.queue.put(request_output)
+                if commit_state:
+                    parent_req = req_state.parent_req
+                    if parent_req is not None:
+                        parent_req.child_requests.discard(request_id)
+                        if not parent_req.child_requests:
+                            self.parent_requests.pop(parent_req.request_id, None)
+            elif parent := self.parent_requests.get(request_id):
+                if parent.child_requests:
+                    child_reqs = list(parent.child_requests)
+                    child_aborted, child_outputs = self.abort_requests_collecting_outputs(
+                        child_reqs,
+                        internal=True,
+                        commit_state=commit_state,
+                    )
+                    request_ids_to_abort.extend(child_aborted)
+                    abort_outputs.extend(child_outputs)
+                if commit_state:
+                    self.parent_requests.pop(request_id, None)
+        return request_ids_to_abort, abort_outputs
+
+    def commit_aborted_request_state(self, request_ids: Iterable[str], *, internal: bool) -> None:
+        """Drop processor mappings after a successful physical EngineCore abort."""
         request_ids = list(request_ids)
         for request_id in request_ids:
             if internal:
@@ -440,7 +551,37 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                     self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
             else:
                 self._native_text_metrics_by_request.pop(request_id, None)
-        return super().abort_requests(request_ids, internal)
+
+        internal_req_ids: list[str] = []
+        for request_id in request_ids:
+            if internal:
+                internal_req_ids.append(request_id)
+                if req_state := self.request_states.get(request_id):
+                    external_req_id = req_state.external_req_id
+                    internal_ids = self.external_req_ids.get(external_req_id)
+                    if internal_ids is not None:
+                        try:
+                            internal_ids.remove(request_id)
+                        except ValueError:
+                            pass
+                        if not internal_ids:
+                            self.external_req_ids.pop(external_req_id, None)
+            elif internal_ids := self.external_req_ids.pop(request_id, []):
+                internal_req_ids.extend(internal_ids)
+
+        for request_id in internal_req_ids:
+            req_state = self.request_states.pop(request_id, None)
+            if req_state is not None:
+                self.lora_states.request_finished(request_id, req_state.lora_name)
+                parent_req = req_state.parent_req
+                if parent_req is not None:
+                    parent_req.child_requests.discard(request_id)
+                    if not parent_req.child_requests:
+                        self.parent_requests.pop(parent_req.request_id, None)
+            elif parent := self.parent_requests.get(request_id):
+                if parent.child_requests:
+                    self.commit_aborted_request_state(list(parent.child_requests), internal=True)
+                self.parent_requests.pop(request_id, None)
 
     def add_request(
         self,

vllm_omni/worker/base.py

@@ -201,11 +201,19 @@ class OmniGPUWorkerBase(GPUWorker):
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
+        # Drain in-flight kernels before CuMem unmap. Abort/decode can still
+        # have GPU work after the scheduler reports idle; unmapping then
+        # returns CUDA_ERROR_INVALID_VALUE (cudaErrorInvalidValue).
+        current_omni_platform.synchronize()
+        gc.collect()
+
         mem_before = current_omni_platform.get_current_memory_usage(self.device)
         offload_tags = ("weights",) if level == 1 else tuple()
         allocator = CuMemAllocator.get_instance()
         allocator.sleep(offload_tags=offload_tags)
-        current_omni_platform.empty_cache()
+        # allocator.sleep already gc.collects and empty_cache()s. A second
+        # empty_cache on the pluggable allocator is a known CUDA invalid
+        # argument (pytorch/pytorch#145168). Only sync for the readout.
         current_omni_platform.synchronize()
         mem_after = current_omni_platform.get_current_memory_usage(self.device)
         freed = max(0, mem_before - mem_after)