[serve][llm] Enable KV cache offloading & KV routing awareness of CPU KV caches

python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py

@@ -61,6 +61,7 @@
 )
 from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
     assign_replica_kv_events_endpoint,
+    enable_native_kv_offload_events,
     get_kv_event_routing_stats,
 )
 from ray.llm._internal.serve.routing_policies.kv_aware.vllm.token_tracking import (
@@ -566,6 +567,8 @@ def _start_async_llm_engine(
         from vllm.v1.executor.abstract import Executor
 
         vllm_engine_config.parallel_config.placement_group = placement_group
+        if is_kv_aware(self.llm_config):
+            enable_native_kv_offload_events(vllm_engine_config)
 
         _clear_current_platform_cache()
 

python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/kv_events.py

@@ -39,20 +39,38 @@ def configure_kv_events_for_kv_routing(llm_config: LLMConfig) -> None:
         }
     )
 
-    _pin_block_hash_seed(llm_config)
+    _configure_runtime_env_for_kv_routing(llm_config)
 
 
-def _pin_block_hash_seed(llm_config: LLMConfig) -> None:
-    """Make engine block hashes content-deterministic across replicas.
+def enable_native_kv_offload_events(vllm_config: Any) -> None:
+    """Make vLLM's native CPU-offload events usable by the KV router."""
+    cache_config = vllm_config.cache_config
+    if (
+        cache_config.kv_offloading_size is None
+        or cache_config.kv_offloading_backend != "native"
+    ):
+        return
+
+    vllm_config.kv_transfer_config.kv_connector_extra_config[
+        "self_describing_kv_events"
+    ] = True
+
+
+def _configure_runtime_env_for_kv_routing(llm_config: LLMConfig) -> None:
+    """Configure vLLM process-wide settings required by KV-aware routing.
 
     The KV router's global indexer chains and dedups blocks by the engines'
     block hashes, so identical content must hash identically on every
     replica. vLLM salts its block-hash chain root per process unless
     ``PYTHONHASHSEED`` is set, so pin it deployment-wide.
+
+    Native offloading must use ``OffloadingConnector`` because it emits the
+    self-describing CPU-tier events consumed by the KV router.
     """
     runtime_env = dict(llm_config.runtime_env or {})
     env_vars = dict(runtime_env.get("env_vars") or {})
     env_vars.setdefault("PYTHONHASHSEED", "0")
+    env_vars["VLLM_USE_SIMPLE_KV_OFFLOAD"] = "0"
     runtime_env["env_vars"] = env_vars
     llm_config.runtime_env = runtime_env
 

python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py

@@ -9,6 +9,7 @@
 from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
     assign_replica_kv_events_endpoint,
     configure_kv_events_for_kv_routing,
+    enable_native_kv_offload_events,
     get_kv_event_routing_stats,
     resolve_kv_event_source_endpoint,
 )
@@ -41,9 +42,16 @@ def ray_instance():
 
 
 class TestConfigureKvEvents:
-    def test_configure_enables_events_and_pins_seed(self):
-        """KV-aware config turns on engine ZMQ KV events and pins the hash seed."""
-        llm_config = make_kv_aware_llm_config()
+    def test_configure_enables_events_and_pins_runtime_env(self):
+        """KV-aware config enables events and required vLLM process settings."""
+        llm_config = make_kv_aware_llm_config(
+            runtime_env={
+                "env_vars": {
+                    "EXISTING_ENV": "value",
+                    "VLLM_USE_SIMPLE_KV_OFFLOAD": "1",
+                }
+            }
+        )
         configure_kv_events_for_kv_routing(llm_config)
 
         assert llm_config.engine_kwargs["kv_events_config"] == {
@@ -53,6 +61,30 @@ def test_configure_enables_events_and_pins_seed(self):
             "replay_endpoint": "tcp://*:6557",
         }
         assert llm_config.runtime_env["env_vars"]["PYTHONHASHSEED"] == "0"
+        assert llm_config.runtime_env["env_vars"]["VLLM_USE_SIMPLE_KV_OFFLOAD"] == "0"
+        assert llm_config.runtime_env["env_vars"]["EXISTING_ENV"] == "value"
+
+    @pytest.mark.parametrize(
+        "offload_size, expected",
+        [
+            (2.0, {"existing": "value", "self_describing_kv_events": True}),
+            (None, {"existing": "value"}),
+        ],
+    )
+    def test_native_offload_event_configuration(self, offload_size, expected):
+        """Only native CPU offload enables complete CPU-tier KV events."""
+        extra_config = {"existing": "value"}
+        vllm_config = SimpleNamespace(
+            cache_config=SimpleNamespace(
+                kv_offloading_size=offload_size,
+                kv_offloading_backend="native",
+            ),
+            kv_transfer_config=SimpleNamespace(kv_connector_extra_config=extra_config),
+        )
+
+        enable_native_kv_offload_events(vllm_config)
+
+        assert extra_config == expected
 
     @pytest.mark.parametrize(
         "engine_kwargs, local_rank, expected_port, expected_replay_port",

release/llm_tests/kv_router_test/test_kv_event_ingestion.py

@@ -54,6 +54,13 @@ async def get_kv_overlap_blocks(self, token_ids: List[int]) -> Dict[int, int]:
 
         Returns the number of leading blocks of ``token_ids`` each worker has cached.
         """
+        scores = await self.get_kv_overlap_scores(token_ids)
+        return {
+            worker_id: score["device_blocks"] for worker_id, score in scores.items()
+        }
+
+    async def get_kv_overlap_scores(self, token_ids: List[int]) -> Dict[int, dict]:
+        """(Test only) Per-worker overlap across every KV storage tier."""
         if self._svc is None:
             return {}
         scores = await self._svc.overlap_scores(
@@ -63,7 +70,7 @@ async def get_kv_overlap_blocks(self, token_ids: List[int]) -> Dict[int, int]:
                 "token_ids": list(token_ids),
             }
         )
-        return {w["worker_id"]: w["device_blocks"] for w in scores["workers"]}
+        return {worker["worker_id"]: worker for worker in scores["workers"]}
 
     async def get_potential_loads(self, token_ids: List[int]) -> Dict[int, dict]:
         """(Test only) Per-worker projected load for routing ``token_ids``."""
@@ -122,7 +129,7 @@ def endpoint(self) -> str:
     def replay_endpoint(self) -> str:
         return f"tcp://127.0.0.1:{self._replay_port}"
 
-    def publish_stored(self, block_hashes, token_ids) -> None:
+    def publish_stored(self, block_hashes, token_ids, medium="GPU") -> None:
         self._pub.publish(
             KVEventBatch(
                 ts=1.0,
@@ -133,18 +140,18 @@ def publish_stored(self, block_hashes, token_ids) -> None:
                         token_ids=list(token_ids),
                         block_size=BLOCK_SIZE,
                         lora_id=None,
-                        medium="GPU",
+                        medium=medium,
                         lora_name=None,
                     )
                 ],
             )
         )
 
-    def publish_removed(self, block_hashes) -> None:
+    def publish_removed(self, block_hashes, medium="GPU") -> None:
         self._pub.publish(
             KVEventBatch(
                 ts=2.0,
-                events=[BlockRemoved(block_hashes=list(block_hashes), medium="GPU")],
+                events=[BlockRemoved(block_hashes=list(block_hashes), medium=medium)],
             )
         )
 
@@ -205,6 +212,23 @@ async def condition():
     await async_wait_for_condition(condition, timeout=timeout, retry_interval_ms=500)
 
 
+async def wait_for_overlap_scores(
+    tracker, token_ids, predicate, publish=None, timeout=30
+):
+    """Poll the tracker's tiered overlap view until ``predicate`` holds."""
+
+    async def condition():
+        if publish is not None:
+            publish()
+        try:
+            scores = await tracker.get_kv_overlap_scores(list(token_ids))
+        except Exception:
+            return False
+        return predicate(scores)
+
+    await async_wait_for_condition(condition, timeout=timeout, retry_interval_ms=500)
+
+
 async def book_uncached_load(tracker, worker_id, count, base):
     """Book ``count`` reservations of distinct, uncached tokens onto ``worker_id``
     so each adds full prefill load to it; returns their reservation ids."""
@@ -388,6 +412,81 @@ async def test_select_worker_prefers_overlap(self):
             a.close()
             b.close()
 
+    @pytest.mark.asyncio
+    async def test_cpu_offload_routing(self):
+        """CPU overlap breaks GPU-overlap ties but loses to a larger GPU hit."""
+        tracker = _LocalKVTokenTracker()
+        a = FakeReplica(23913)
+        b = FakeReplica(23914)
+        worker_a = get_worker_id("replica-A")
+        worker_b = get_worker_id("replica-B")
+        block_hashes = [701, 702, 703]
+        token_ids = list(range(3 * BLOCK_SIZE))
+        try:
+            tracker._on_deployment_targets(
+                targets(
+                    running_replica("replica-A", a.endpoint(), a.replay_endpoint()),
+                    running_replica("replica-B", b.endpoint(), b.replay_endpoint()),
+                )
+            )
+            await wait_registered(tracker, [worker_a, worker_b])
+
+            a.publish_stored(block_hashes, token_ids)
+            a.publish_stored(block_hashes, token_ids, medium="CPU")
+            a.publish_removed(block_hashes[1:])
+            b.publish_stored(block_hashes[:1], token_ids[:BLOCK_SIZE])
+            await wait_for_overlap_scores(
+                tracker,
+                token_ids,
+                lambda scores: (
+                    scores.get(worker_a, {}).get("device_blocks") == 1
+                    and scores.get(worker_a, {}).get("host_pinned_extension_blocks")
+                    == 2
+                    and scores.get(worker_b, {}).get("device_blocks") == 1
+                ),
+            )
+
+            cpu_selection = await tracker.select_worker(
+                "cpu-tiebreak", token_ids, [worker_a, worker_b]
+            )
+            assert cpu_selection["worker_id"] == worker_a
+            await tracker.on_request_completed("cpu-tiebreak")
+
+            b.publish_stored(block_hashes, token_ids)
+            await wait_for_overlap_scores(
+                tracker,
+                token_ids,
+                lambda scores: scores.get(worker_b, {}).get("device_blocks") == 3,
+            )
+
+            scores = await tracker.get_kv_overlap_scores(token_ids)
+            assert scores[worker_a]["host_pinned_extension_blocks"] == 2
+            assert 1 < scores[worker_a]["router_credit_blocks"] < 3
+            assert scores[worker_b]["router_credit_blocks"] == 3
+
+            gpu_selection = await tracker.select_worker(
+                "gpu-hit", token_ids, [worker_a, worker_b]
+            )
+            assert gpu_selection["worker_id"] == worker_b
+            assert gpu_selection["effective_prefill_tokens"] == 0
+            await tracker.on_request_completed("gpu-hit")
+
+            b.publish_removed(block_hashes)
+            await wait_for_overlap(
+                tracker,
+                token_ids,
+                lambda overlap: overlap.get(worker_b) == 0,
+            )
+
+            cpu_selection = await tracker.select_worker(
+                "cpu-reload", token_ids, [worker_a, worker_b]
+            )
+            assert cpu_selection["worker_id"] == worker_a
+            assert 0 < cpu_selection["effective_prefill_tokens"] < len(token_ids)
+        finally:
+            a.close()
+            b.close()
+
     @pytest.mark.asyncio
     async def test_select_worker_prefers_lower_load(self):
         """With equal KV overlap on both workers, select_worker routes to the