[Rebase] Align vLLM-Omni with vLLM v0.31.0 (#8459)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/entrypoints/test_async_omni.py

@@ -10,6 +10,7 @@ import pytest
 from vllm.engine.protocol import StreamingInput
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.v1.kv_hints import KvHintsEnvelope
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.stage_config import get_deploy_config_path
@@ -138,6 +139,39 @@ def test_generate_forwards_lora_request_to_engine():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("streaming", [None, "chunk", "empty"])
+def test_generate_forwards_kv_hints_on_initial_stage_submission(streaming):
+    async def run():
+        submissions = []
+
+        async def add(**kwargs):
+            submissions.append(kwargs)
+
+        omni = get_async_omni_instance(fake_add_request=add)
+        omni.engine.add_streaming_update_async = _noop
+        omni.engine.stage_configs = [SimpleNamespace(is_comprehension=True, stage_type="llm")]
+        omni.engine.stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))]
+        hints = KvHintsEnvelope("1", "message", [])
+        params = SamplingParams(max_tokens=2, output_kind=RequestOutputKind.DELTA)
+
+        async def chunks():
+            if streaming == "chunk":
+                yield StreamingInput(prompt={"prompt_token_ids": [1]})
+
+        async for _ in omni.generate(
+            prompt=chunks() if streaming else {"prompt_token_ids": [1]},
+            request_id="hints",
+            sampling_params_list=[params],
+            kv_hints=hints,
+        ):
+            pass
+        assert len(submissions) == 1
+        assert submissions[0]["kv_hints"] is hints
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize(
     "req_ids,cancel_prefix,expected_cancel_count",

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -194,6 +194,56 @@ def test_reset_encoder_cache_forwards_to_ar_stages():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_release_kv_memory_requires_pause_and_wakes_only_ar_kv_allocations():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+        with pytest.raises(RuntimeError, match="completed pause"):
+            await omni.release_kv_cache_memory()
+        omni.collective_rpc.assert_not_awaited()
+
+        omni._paused = True
+        await omni.release_kv_cache_memory()
+        omni._clear_frontend_mm_cache.assert_awaited_once_with()
+        omni.collective_rpc.assert_awaited_once_with(
+            method="release_kv_cache_memory",
+            args=(),
+            kwargs=None,
+            stage_ids=[0, 2],
+            timeout=CACHE_RESET_TIMEOUT_S,
+        )
+        assert omni._stage_sleeping_tags == {0: {"kv_cache"}, 2: {"kv_cache"}}
+        assert not omni._level2_sleeping
+
+        omni.collective_rpc.reset_mock()
+        await omni.wake_up(stage_ids=[0, 2], tags=["kv_cache"])
+        omni.collective_rpc.assert_awaited_once_with(
+            method="wake_up", args=(), kwargs={"tags": ["kv_cache"]}, stage_ids=[0, 2]
+        )
+        assert not omni._sleeping_tags
+        assert omni._paused
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_release_kv_memory_skips_diffusion_and_reports_engine_failures():
+    async def run() -> None:
+        diffusion = _make_omni(stage_types=["diffusion"])
+        await diffusion.release_kv_cache_memory()
+        diffusion.collective_rpc.assert_not_awaited()
+        assert not diffusion._paused
+
+        omni = _make_omni(stage_types=["llm"])
+        omni._paused = True
+        omni.collective_rpc.return_value = [{"supported": False, "error": "requires all executor memory resident"}]
+        with pytest.raises(RuntimeError, match="all executor memory resident"):
+            await omni.release_kv_cache_memory()
+        assert not omni._sleeping_tags
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages():
     async def run() -> None: