[Orchestrator] Implement cache reset RPC forwarding to AR stage engine cores (#6766)

Signed-off-by: Asthenia <asthenia0412@gmail.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Asthenia <asthenia0412@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/engine/test_orchestrator_error_handling.py

@@ -1206,3 +1206,110 @@ async def test_rpc_failure_capture_is_limited_to_pause_and_resume() -> None:
     finally:
         for q in queues:
             q.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["reset_mm_cache", "reset_encoder_cache", "reset_prefix_cache"])
+async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(method):
+    clients = [FakeStageClient(stage_type="llm") for _ in range(2)]
+    setattr(clients[0], f"{method}_async", AsyncMock(side_effect=RuntimeError("reset failed")))
+    ok = AsyncMock(return_value=True)
+    setattr(clients[1], f"{method}_async", ok)
+    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([clients]))
+    try:
+        await orchestrator._handle_collective_rpc(
+            CollectiveRPCRequestMessage(rpc_id="reset", method=method, args=(), kwargs={}, stage_ids=[0], timeout=1.0)
+        )
+        result = queues[2].async_q.get_nowait()
+        assert result.results[0]["supported"] is False
+        assert "reset failed" in result.results[0]["error"]
+        assert result.results[1] is True
+        ok.assert_awaited_once_with()
+        # Errors before StagePool dispatch (e.g. a missing route) must also
+        # be serialized at the orchestrator boundary, then allow another RPC.
+        orchestrator.stage_pools[0].collective_rpc = AsyncMock(side_effect=RuntimeError("route failed"))
+        await orchestrator._handle_collective_rpc(
+            CollectiveRPCRequestMessage(rpc_id="route", method=method, args=(), kwargs={}, stage_ids=[0])
+        )
+        assert all("route failed" in r["error"] for r in queues[2].async_q.get_nowait().results)
+        orchestrator.stage_pools[0].collective_rpc = AsyncMock(return_value=True)
+        await orchestrator._handle_collective_rpc(
+            CollectiveRPCRequestMessage(rpc_id="next", method=method, args=(), kwargs={}, stage_ids=[0])
+        )
+        assert queues[2].async_q.get_nowait().results == [True, True]
+    finally:
+        for q in queues:
+            q.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "method, kwargs, clears",
+    [
+        ("reset_mm_cache", {}, True),
+        ("sleep", {}, True),
+        ("pause_scheduler", {"clear_cache": True}, True),
+        ("pause_scheduler", {"clear_cache": False}, False),
+        ("reset_prefix_cache", {}, False),
+        ("reset_encoder_cache", {}, False),
+    ],
+)
+async def test_control_rpc_clears_only_selected_downstream_sender(method, kwargs, clears):
+    pools = _build_stage_pools([[FakeStageClient(stage_type="llm") for _ in range(2)] for _ in range(2)])
+    orchestrator, queues = _build_bare_orchestrator(pools)
+    events = []
+    renderers = [SimpleNamespace(clear_mm_cache_async=AsyncMock()) for _ in range(2)]
+    renderers[1].clear_mm_cache_async.side_effect = lambda: events.append("sender")
+    orchestrator._stage_input_processors = {i: SimpleNamespace(renderer=r) for i, r in enumerate(renderers)}
+
+    async def rpc(**kwargs):
+        events.append("receiver")
+        return True
+
+    pools[0].collective_rpc = AsyncMock()
+    pools[1].collective_rpc = AsyncMock(side_effect=rpc)
+    try:
+        await orchestrator._handle_collective_rpc(
+            CollectiveRPCRequestMessage(
+                rpc_id="reset", method=method, args=(), kwargs=kwargs, stage_ids=[1], timeout=1.0
+            )
+        )
+        renderers[0].clear_mm_cache_async.assert_not_awaited()
+        pools[0].collective_rpc.assert_not_awaited()
+        assert renderers[1].clear_mm_cache_async.await_count == int(clears)
+        assert events == (["sender"] if clears else []) + ["receiver", "receiver"]
+        assert queues[2].async_q.get_nowait().results == [True, True]
+    finally:
+        for q in queues:
+            q.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", ["error", "timeout"])
+async def test_downstream_sender_clear_failure_is_nonfatal(failure):
+    pools = _build_stage_pools([[FakeStageClient(stage_type="llm")]])
+    orchestrator, queues = _build_bare_orchestrator(pools)
+
+    async def clear():
+        if failure == "timeout":
+            await asyncio.Event().wait()
+        raise RuntimeError("sender cache failed")
+
+    renderer = SimpleNamespace(clear_mm_cache_async=AsyncMock(side_effect=clear))
+    orchestrator._stage_input_processors[0] = SimpleNamespace(renderer=renderer)
+    pools[0].collective_rpc = AsyncMock(return_value=None)
+    message = CollectiveRPCRequestMessage(
+        rpc_id="reset", method="reset_mm_cache", args=(), kwargs={}, stage_ids=[0], timeout=0.01
+    )
+    try:
+        await orchestrator._handle_collective_rpc(message)
+        result = queues[2].async_q.get_nowait().results[0]
+        assert result["supported"] is False
+        assert ("TimeoutError" if failure == "timeout" else "sender cache failed") in result["error"]
+        pools[0].collective_rpc.assert_not_awaited()
+        renderer.clear_mm_cache_async = AsyncMock()
+        await orchestrator._handle_collective_rpc(message)
+        assert queues[2].async_q.get_nowait().results == [None]
+    finally:
+        for q in queues:
+            q.close()

tests/engine/test_stage_pool_collective_rpc.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unit tests for StagePool.collective_rpc EngineCore control dispatch."""
 
 from __future__ import annotations
@@ -40,14 +43,14 @@ def test_collective_rpc_normalizes_none_args_on_control_helper():
 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)
+        pool, client = _make_pool(is_sleeping_async=other)
 
-        result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1})
+        result = await pool.collective_rpc(0, "is_sleeping", 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",
+            method="is_sleeping",
             timeout=1.5,
             args=("x",),
             kwargs={"k": 1},
@@ -63,7 +66,7 @@ def test_collective_rpc_control_helper_honors_timeout():
             await asyncio.sleep(1.0)
             return "slept"
 
-        pool, _client = _make_pool(sleep_async=slow_sleep)
+        pool, _client = _make_pool(sleep_async=AsyncMock(side_effect=slow_sleep))
         with pytest.raises(asyncio.TimeoutError):
             await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,))
 
@@ -80,11 +83,38 @@ def test_collective_rpc_control_method_reraises_worker_error():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_collective_rpc_reset_caches_use_control_helpers():
+    async def run() -> None:
+        reset_prefix = AsyncMock(return_value=True)
+        reset_encoder = AsyncMock(return_value=None)
+        reset_mm = AsyncMock(return_value=None)
+        pool, client = _make_pool(
+            reset_prefix_cache_async=reset_prefix,
+            reset_encoder_cache_async=reset_encoder,
+            reset_mm_cache_async=reset_mm,
+        )
+
+        result = await pool.collective_rpc(0, "reset_prefix_cache", args=(True, True), timeout=2.0)
+        assert result is True
+        reset_prefix.assert_awaited_once_with(True, True)
+
+        await pool.collective_rpc(0, "reset_encoder_cache", timeout=2.0)
+        reset_encoder.assert_awaited_once()
+
+        await pool.collective_rpc(0, "reset_mm_cache", timeout=2.0)
+        reset_mm.assert_awaited_once()
+
+        client.collective_rpc_async.assert_not_awaited()
+
+    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")
+        result = await pool.collective_rpc(0, "is_sleeping")
         assert result["supported"] is False
         assert "probe failed" in result["error"]
 
@@ -122,3 +152,28 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
         abort.assert_awaited_once()
 
     asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
+@pytest.mark.parametrize("failure", ["error", "timeout", "missing"])
+def test_cache_reset_failure_is_serialized_and_pool_remains_usable(method, failure):
+    async def run():
+        async def reset():
+            if failure == "timeout":
+                await asyncio.Event().wait()
+            raise RuntimeError("cache reset failed")
+
+        pool, client = _make_pool()
+        if failure != "missing":
+            setattr(client, f"{method}_async", AsyncMock(side_effect=reset))
+        result = await pool.collective_rpc(0, method, timeout=0.01)
+        assert result["supported"] is False
+        assert result["error"]
+        if failure == "timeout":
+            assert "timed out" in result["error"]
+        client.collective_rpc_async.assert_not_awaited()
+        # A subsequent control request still completes on the same pool.
+        assert await pool.collective_rpc(0, "is_sleeping") == {"via": "collective"}
+
+    asyncio.run(run())

tests/entrypoints/test_async_omni_cache_reset.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Real EngineCore cache invalidation through the AsyncOmni control plane."""
+
+import pytest
+from PIL import Image
+from vllm import SamplingParams
+
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.cuda
+@pytest.mark.asyncio
+async def test_cache_resets_invalidate_real_engine_and_allow_repeated_multimodal_inputs():
+    deploy_config = modify_stage_config(
+        get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enable_prefix_caching": True,
+                    "mm_processor_cache_gb": 0.1,
+                    "max_model_len": 2048,
+                    "max_num_batched_tokens": 2048,
+                    "num_gpu_blocks_override": 256,
+                    "enforce_eager": True,
+                }
+            }
+        },
+    )
+    engine = AsyncOmni(
+        model="Qwen/Qwen2.5-Omni-7B",
+        deploy_config=deploy_config,
+        stage_init_timeout=600,
+    )
+
+    async def generate(prompt, request_id):
+        final = None
+        async for output in engine.generate(
+            prompt,
+            request_id=request_id,
+            sampling_params_list=[SamplingParams(max_tokens=4, temperature=0)],
+        ):
+            final = output
+        assert final is not None and final.finished
+        assert final.outputs and final.outputs[0].token_ids
+        return final
+
+    try:
+        prompt = "Explain how a cache works. " * 40
+        cold = await generate(prompt, "cold")
+        warm = await generate(prompt, "warm")
+        assert cold.num_cached_tokens == 0
+        assert warm.num_cached_tokens > 0
+
+        assert await engine.reset_prefix_cache(reset_running_requests=True, reset_connector=True)
+        await engine.reset_encoder_cache()
+        await engine.reset_mm_cache()
+        reset = await generate(prompt, "reset")
+        assert reset.num_cached_tokens == 0
+        assert reset.outputs[0].token_ids == cold.outputs[0].token_ids
+
+        image_prompt = {
+            "prompt": "<|im_start|>user\n<|vision_start|><|IMAGE|><|vision_end|>Describe this image.<|im_end|>\n<|im_start|>assistant\n",
+            "multi_modal_data": {"image": Image.new("RGB", (56, 56), "red")},
+        }
+        before = await generate(image_prompt, "image-before")
+        await engine.reset_mm_cache()
+        await engine.reset_encoder_cache()
+        assert await engine.reset_prefix_cache()
+        after = await generate(image_prompt, "image-after")
+        assert after.outputs[0].token_ids == before.outputs[0].token_ids
+    finally:
+        engine.shutdown()

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -14,7 +14,7 @@ from vllm.renderers import BaseRenderer
 from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.diffusion.data import CuMemTag, OmniACK
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.async_omni import CACHE_RESET_TIMEOUT_S, AsyncOmni
 
 pytestmark = [pytest.mark.core_model]
 
@@ -31,7 +31,7 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
     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()
+    omni._clear_frontend_mm_cache = AsyncMock()
 
     stage_configs = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
     omni.engine = SimpleNamespace(
@@ -46,7 +46,8 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
 @pytest.mark.cpu
 def test_reset_mm_cache_routes_through_renderer(mocker):
     async def run() -> None:
-        omni = object.__new__(AsyncOmni)
+        omni = _make_omni(stage_types=["llm"])
+        del omni._clear_frontend_mm_cache
         renderer = mocker.create_autospec(BaseRenderer, instance=True)
         input_processor = mocker.create_autospec(InputProcessor, instance=True)
         input_processor.renderer = renderer
@@ -115,9 +116,10 @@ def test_pause_generation_still_rpcs_when_already_paused():
             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()
+        omni.reset_prefix_cache.assert_not_awaited()
+        omni.reset_mm_cache.assert_not_awaited()
+        omni.reset_encoder_cache.assert_not_awaited()
+        omni._clear_frontend_mm_cache.assert_awaited_once_with()
 
     asyncio.run(run())
 
@@ -142,6 +144,84 @@ def test_resume_generation_resumes_ar_then_clears_frontend_pause():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_reset_prefix_cache_forwards_to_ar_stages():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+
+        result = await omni.reset_prefix_cache(reset_running_requests=True, reset_connector=True)
+
+        assert result is True
+        omni.collective_rpc.assert_awaited_once_with(
+            method="reset_prefix_cache",
+            args=(True, True),
+            kwargs=None,
+            stage_ids=[0],
+            timeout=CACHE_RESET_TIMEOUT_S,
+        )
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_reset_prefix_cache_diffusion_only_skips_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+
+        result = await omni.reset_prefix_cache()
+
+        assert result is True
+        omni.collective_rpc.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_reset_encoder_cache_forwards_to_ar_stages():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+
+        await omni.reset_encoder_cache()
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="reset_encoder_cache",
+            args=(),
+            kwargs=None,
+            stage_ids=[0],
+            timeout=CACHE_RESET_TIMEOUT_S,
+        )
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        calls = []
+        renderer = SimpleNamespace(clear_mm_cache_async=AsyncMock(side_effect=lambda: calls.append("frontend")))
+        omni.input_processor = SimpleNamespace(renderer=renderer)
+        del omni._clear_frontend_mm_cache
+
+        async def rpc(**kwargs):
+            calls.append("engine")
+            return [None]
+
+        omni.collective_rpc.side_effect = rpc
+        await omni.reset_mm_cache()
+
+        assert calls == ["frontend", "engine"]
+        omni.collective_rpc.assert_awaited_once_with(
+            method="reset_mm_cache",
+            args=(),
+            kwargs=None,
+            stage_ids=[0],
+            timeout=CACHE_RESET_TIMEOUT_S,
+        )
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
     async def run() -> None:
@@ -539,8 +619,9 @@ def test_mixed_engine_keep_sends_engine_core_and_diffusion_pauses():
             "kwargs": {"mode": "keep"},
             "stage_ids": [1],
         }
-        # The frontend cache reset still runs after the backend pauses.
-        omni.reset_prefix_cache.assert_awaited_once()
+        # Only P0 needs clearing here: pause_scheduler already resets P1.
+        omni.reset_prefix_cache.assert_not_awaited()
+        omni._clear_frontend_mm_cache.assert_awaited_once_with()
 
         await omni.resume_generation()
         assert _rpc_methods(omni)[2:] == [("resume_scheduler", [0]), ("resume_scheduler", [1])]
@@ -579,9 +660,11 @@ def test_keep_then_sleep_wake_resume_keeps_admission_paused_until_resume():
         omni.reset_prefix_cache = AsyncMock(return_value=True)
         omni.reset_encoder_cache = AsyncMock()
         omni.collective_rpc = AsyncMock(
-            side_effect=lambda **kw: [OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)]
-            if kw["method"] in {"handle_sleep_task", "handle_wake_task"}
-            else [None]
+            side_effect=lambda **kw: (
+                [OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)]
+                if kw["method"] in {"handle_sleep_task", "handle_wake_task"}
+                else [None]
+            )
         )
 
         await omni.pause_generation(mode="keep", clear_cache=False)
@@ -683,3 +766,63 @@ def test_two_ar_stages_partial_resume_keeps_admission_paused():
         assert await omni.is_paused() is False
 
     asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
+@pytest.mark.parametrize("stage_ids, targets", [(None, [0, 2]), ([2], [2]), ([1], []), ([], [])])
+def test_cache_resets_respect_stage_selection(method, stage_ids, targets):
+    async def run():
+        omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+        await getattr(omni, method)(stage_ids=stage_ids, timeout=2.0)
+        if targets:
+            assert omni.collective_rpc.await_args.kwargs["stage_ids"] == targets
+            assert omni.collective_rpc.await_args.kwargs["timeout"] == 2.0
+        else:
+            omni.collective_rpc.assert_not_awaited()
+        assert omni._clear_frontend_mm_cache.await_count == int(method == "reset_mm_cache" and 0 in targets)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
+@pytest.mark.parametrize("failure", [{"todo": "unsupported"}, {"supported": False}, {"error": "reset failed"}])
+def test_cache_resets_reject_failure_from_any_replica(method, failure):
+    async def run():
+        omni = _make_omni(stage_types=["llm"])
+        # Exercise the actual collective_rpc too: a single stage can have
+        # multiple replicas, so result indexes are not stage IDs.
+        del omni.collective_rpc
+        omni.engine.collective_rpc_async.return_value = [True, failure]
+        with pytest.raises(RuntimeError, match=method):
+            await getattr(omni, method)()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_prefix_reset_returns_false_if_any_replica_refuses():
+    async def run():
+        omni = _make_omni(stage_types=["llm"])
+        omni.collective_rpc.return_value = [True, [True, False]]
+        assert await omni.reset_prefix_cache() is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("operation", ["pause_generation", "sleep"])
+@pytest.mark.parametrize("stage_ids", [[0], [2]])
+def test_pause_and_sleep_do_not_reset_unselected_stages_or_reset_twice(operation, stage_ids):
+    async def run():
+        omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+        kwargs = {"clear_cache": True} if operation == "pause_generation" else {}
+        await getattr(omni, operation)(stage_ids=stage_ids, **kwargs)
+        assert omni.collective_rpc.await_count == 1
+        rpc = omni.collective_rpc.await_args.kwargs
+        assert rpc["method"] == ("pause_scheduler" if operation == "pause_generation" else "sleep")
+        assert rpc["stage_ids"] == stage_ids
+        assert omni._clear_frontend_mm_cache.await_count == int(0 in stage_ids)
+
+    asyncio.run(run())

vllm_omni/engine/orchestrator.py

@@ -704,8 +704,21 @@ class OrchestratorBase:
         results: list[Any] = []
         stage_ids: list[int] = []
         for pool in target_pools:
+            clear_sender = pool.stage_type != "diffusion" and (
+                method in ("reset_mm_cache", "sleep")
+                or (method == "pause_scheduler" and kwargs.get("clear_cache", args[1] if len(args) > 1 else True))
+            )
             for replica_id in pool.live_replica_ids():
                 try:
+                    if clear_sender:
+                        processor = self._stage_input_processors.get(pool.stage_id)
+                        if processor is not None:
+                            clear = processor.renderer.clear_mm_cache_async()
+                            if timeout is None:
+                                await clear
+                            else:
+                                await asyncio.wait_for(clear, timeout=timeout)
+                        clear_sender = False
                     stage_result = await pool.collective_rpc(
                         replica_id=replica_id,
                         method=method,
@@ -714,13 +727,14 @@ class OrchestratorBase:
                         kwargs=kwargs,
                     )
                 except Exception as exc:
-                    if method not in ("pause_scheduler", "resume_scheduler"):
+                    if (
+                        method not in ("pause_scheduler", "resume_scheduler")
+                        and method not in StagePool._CACHE_RESET_METHODS
+                    ):
                         raise
-                    # A pause or resume that fails on one replica leaves the
-                    # engine usable, so report it as this replica's result
-                    # rather than out of the request handler; the caller
-                    # reaches both through _engine_core_rpc, which raises on
-                    # the error result and can then retry or resume.
+                    # Administrative failures must reach the caller without
+                    # terminating unrelated stages. _engine_core_rpc raises
+                    # on this result; remaining replicas still receive the RPC.
                     stage_result = {"supported": False, "error": f"{type(exc).__name__}: {exc}"}
                 stage_ids.append(pool.stage_id)
                 results.append(stage_result)

vllm_omni/engine/stage_pool.py

@@ -101,12 +101,16 @@ class StagePool:
     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.
+    _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"})
     _ENGINE_CORE_CONTROL_ASYNC_METHODS = frozenset(
         {
             "pause_scheduler",
             "resume_scheduler",
             "sleep",
             "wake_up",
+            "reset_prefix_cache",
+            "reset_encoder_cache",
+            "reset_mm_cache",
         }
     )
 
@@ -1330,6 +1334,8 @@ class StagePool:
                     if timeout is not None:
                         return await asyncio.wait_for(result, timeout=timeout)
                     return await result
+                if method in self._CACHE_RESET_METHODS:
+                    return {"supported": False, "error": f"EngineCore helper {method}_async is unavailable"}
 
             return await client.collective_rpc_async(
                 method=method,
@@ -1344,7 +1350,7 @@ class StagePool:
                 replica_id,
                 method,
             )
-            if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS:
+            if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS and method not in self._CACHE_RESET_METHODS:
                 raise
             if isinstance(exc, TimeoutError):
                 error = f"{type(exc).__name__}: {method} timed out after {timeout}s"

vllm_omni/entrypoints/async_omni.py

@@ -46,6 +46,8 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+CACHE_RESET_TIMEOUT_S = 60.0
+
 
 class AsyncOmni(AsyncOmniBase, EngineClient):
     """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine.
@@ -535,17 +537,15 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             stage_ids=stage_ids,
         )
 
-        unsupported_stage_ids: list[int] = []
-        effective_stage_ids = stage_ids or list(range(len(results)))
-        for index, result in enumerate(results):
-            if isinstance(result, dict) and result.get("todo"):
-                unsupported_stage_ids.append(effective_stage_ids[index])
-
-        if unsupported_stage_ids:
+        unsupported_results = [
+            index for index, result in enumerate(results) if isinstance(result, dict) and result.get("todo")
+        ]
+        if unsupported_results:
             logger.warning(
-                "[AsyncOmni] collective_rpc(%s) has TODO support on stage(s): %s",
+                "[AsyncOmni] collective_rpc(%s) has TODO support in replica result(s) %s (requested stages: %s)",
                 method,
-                unsupported_stage_ids,
+                unsupported_results,
+                stage_ids,
             )
 
         return results
@@ -557,23 +557,28 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         stage_ids: list[int],
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
+        timeout: float | None = None,
     ) -> list[Any]:
-        """Call an engine control helper via collective_rpc (orchestrator loop).
-
-        StagePool resolves ``{method}_async`` on the AR client when present
-        (vLLM AsyncMPClient convention); diffusion stages answer the same
-        method names inside DiffusionEngine. Raises if any replica reports
-        failure.
-        """
+        """Call EngineCore helpers and reject failures from any replica."""
+        timeout_kwargs = {"timeout": timeout} if timeout is not None else {}
         results = await self.collective_rpc(
             method=method,
             args=args,
             kwargs=kwargs,
             stage_ids=stage_ids,
+            **timeout_kwargs,
         )
+
+        def check_result(result: Any) -> None:
+            if isinstance(result, (list, tuple)):
+                for item in result:
+                    check_result(item)
+            elif isinstance(result, dict):
+                if result.get("error") or result.get("todo") or result.get("supported") is False:
+                    raise RuntimeError(f"{method} failed: {result}")
+
         for result in results:
-            if isinstance(result, dict) and result.get("error"):
-                raise RuntimeError(f"{method} failed: {result['error']}")
+            check_result(result)
         return results
 
     @staticmethod
@@ -768,13 +773,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
 
         # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler
         # already clears AR-side caches when clear_cache=True.
-        if clear_cache:
-            await self.reset_prefix_cache(
-                reset_running_requests=not wait_for_inflight_requests,
-                reset_connector=True,
-            )
-            await self.reset_mm_cache()
-            await self.reset_encoder_cache()
+        if clear_cache and 0 in ar_stage_ids:
+            await self._clear_frontend_mm_cache()
 
     async def resume_generation(self, stage_ids: list[int] | None = None) -> None:
         """Resume generation after :meth:`pause_generation`."""
@@ -835,34 +835,63 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         """
         return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages)
 
-    async def reset_mm_cache(self) -> None:
-        """Reset the frontend (P0) multimodal processor cache.
-
-        ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache.
-        Clearing P0 avoids hash-only follow-up requests after that reset.
-        """
+    async def _clear_frontend_mm_cache(self) -> None:
+        """Clear P0 through the renderer's serialized multimodal executor."""
         renderer = self.renderer
         if renderer is not None:
             await renderer.clear_mm_cache_async()
 
-    async def reset_encoder_cache(self) -> None:
-        """Reset the encoder cache for all stages.
+    async def reset_mm_cache(
+        self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S
+    ) -> None:
+        """Reset sender and receiver MM caches on selected AR stages.
 
-        TODO: Forward to Orchestrator process via message.
+        By default all AR stages are selected; diffusion stages are skipped.
+        Stage 0's sender is the frontend renderer. Downstream senders are
+        cleared by the orchestrator before resetting their engine cores.
+        Call while generation is paused to avoid racing new inputs.
         """
-        logger.warning("[AsyncOmni] reset_encoder_cache not yet supported with Orchestrator process")
+        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
+        if 0 in ar_stage_ids:
+            await asyncio.wait_for(self._clear_frontend_mm_cache(), timeout=timeout)
+        if ar_stage_ids:
+            await self._engine_core_rpc("reset_mm_cache", stage_ids=ar_stage_ids, timeout=timeout)
+
+    async def reset_encoder_cache(
+        self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S
+    ) -> None:
+        """Reset encoder caches on selected AR stages (all by default).
+
+        Diffusion stages are skipped. RPC failures are raised to the caller.
+        """
+        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
+        if ar_stage_ids:
+            await self._engine_core_rpc("reset_encoder_cache", stage_ids=ar_stage_ids, timeout=timeout)
 
     async def reset_prefix_cache(
         self,
         reset_running_requests: bool = False,
         reset_connector: bool = False,
+        *,
+        stage_ids: list[int] | None = None,
+        timeout: float = CACHE_RESET_TIMEOUT_S,
     ) -> bool:
-        """Reset the prefix cache for all stages.
+        """Reset prefix caches on selected AR stages (all by default).
 
-        TODO: Forward to Orchestrator process via message.
+        Diffusion stages are skipped. Return False if a stage cannot reset
+        its cache; unsupported operations and RPC failures raise instead of
+        silently retaining KV computed under previous model weights.
         """
-        logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process")
-        return True
+        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
+        if not ar_stage_ids:
+            return True
+        results = await self._engine_core_rpc(
+            "reset_prefix_cache",
+            stage_ids=ar_stage_ids,
+            args=(reset_running_requests, reset_connector),
+            timeout=timeout,
+        )
+        return all(self._coerce_stage_bool(result) for result in results)
 
     async def sleep(
         self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort"
@@ -895,11 +924,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             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()
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        # EngineCore.sleep resets receiver caches itself; only clear P0 here.
+        if 0 in ar_stage_ids:
+            await self._clear_frontend_mm_cache()
 
         self._final_output_handler()
-        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