From f82e4e4861bd447cc26475957156243fe9a91b17 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Mon, 27 Jul 2026 17:19:33 -0700 Subject: [PATCH 1/4] [serve][llm] Enable KV cache offloading & KV routing awareness of CPU KV caches Signed-off-by: Jeffrey Wang --- .../serve/engines/vllm/vllm_engine.py | 3 + .../kv_aware/vllm/kv_events.py | 24 ++- .../deployments/routers/kv/test_kv_events.py | 48 ++++- .../kv_router_test/test_kv_event_ingestion.py | 109 +++++++++++- .../kv_router_test/test_kv_events.py | 164 ++++++++++++++++++ release/llm_tests/kv_router_test/utils.py | 9 +- 6 files changed, 345 insertions(+), 12 deletions(-) diff --git a/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py b/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py index 83a1e4170708..00278784fc07 100644 --- a/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py +++ b/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py @@ -58,6 +58,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 ( @@ -564,6 +565,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() diff --git a/python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/kv_events.py b/python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/kv_events.py index 649d45320efb..6f88f2942b5c 100644 --- a/python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/kv_events.py +++ b/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 diff --git a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py index 46a2b92aa6ca..4d12979be357 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py +++ b/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,40 @@ 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" + + def test_native_offload_enables_self_describing_events(self): + """The resolved native connector emits complete CPU-tier KV events.""" + extra_config = {"offload_prompt_only": False} + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_offloading_size=2.0, + kv_offloading_backend="native", + ), + kv_transfer_config=SimpleNamespace(kv_connector_extra_config=extra_config), + ) + + enable_native_kv_offload_events(vllm_config) + + assert extra_config == { + "offload_prompt_only": False, + "self_describing_kv_events": True, + } + + def test_no_offload_is_unchanged(self): + extra_config = {"existing": "value"} + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_offloading_size=None, + kv_offloading_backend="native", + ), + kv_transfer_config=SimpleNamespace(kv_connector_extra_config=extra_config), + ) + + enable_native_kv_offload_events(vllm_config) + + assert extra_config == {"existing": "value"} @pytest.mark.parametrize( "engine_kwargs, local_rank, expected_port, expected_replay_port", diff --git a/release/llm_tests/kv_router_test/test_kv_event_ingestion.py b/release/llm_tests/kv_router_test/test_kv_event_ingestion.py index ac4f0e45d53e..b1e32638722d 100644 --- a/release/llm_tests/kv_router_test/test_kv_event_ingestion.py +++ b/release/llm_tests/kv_router_test/test_kv_event_ingestion.py @@ -52,6 +52,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( @@ -61,7 +68,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``.""" @@ -120,7 +127,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, @@ -131,18 +138,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)], ) ) @@ -203,6 +210,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.""" @@ -395,6 +419,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 diff --git a/release/llm_tests/kv_router_test/test_kv_events.py b/release/llm_tests/kv_router_test/test_kv_events.py index 749680fa620d..f83d51ac4756 100644 --- a/release/llm_tests/kv_router_test/test_kv_events.py +++ b/release/llm_tests/kv_router_test/test_kv_events.py @@ -44,6 +44,16 @@ FLUSH_MESSAGES = [ {"role": "user", "content": _SHARED_PREFIX + " beside the tall fence."} ] +OFFLOAD_MESSAGES = [ + { + "role": "user", + "content": ( + "Remember this native KV offload target exactly. " + + ("alpha beta gamma delta epsilon zeta eta theta iota kappa " "lambda mu ") + * 20 + ), + } +] def post_chat(endpoint, messages=MESSAGES, max_tokens=MAX_TOKENS): @@ -296,6 +306,160 @@ async def manual_fully_overlaps(): assert chat_token_ids == manual_token_ids +class TestKvOffload: + """End-to-end native vLLM CPU offload with tier-aware routing.""" + + @pytest.fixture(scope="class") + def deployed_handle(self): + if not ray.is_initialized(): + ray.init(address="auto") + serve.shutdown() + + llm_config = LLMConfig( + model_loading_config=ModelLoadingConfig( + model_id=MODEL_ID, + model_source=MODEL_SOURCE, + ), + deployment_config=dict( + autoscaling_config=dict( + min_replicas=NUM_REPLICAS, max_replicas=NUM_REPLICAS + ), + request_router_config=RequestRouterConfig( + request_router_class=KVAwareRouter + ), + ), + engine_kwargs=dict( + enable_prefix_caching=True, + enable_prompt_tokens_details=True, + enable_force_include_usage=True, + enforce_eager=True, + gpu_memory_utilization=0.4, + kv_offloading_backend="native", + kv_offloading_size=1.0, + max_model_len=512, + num_gpu_blocks_override=32, + ), + experimental_configs={"KV_EVENTS_PORT_BASE": 21700}, + runtime_env=dict( + env_vars={ + "RAY_SERVE_ENABLE_DIRECT_INGRESS": "1", + "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING": "1", + } + ), + log_engine_metrics=False, + ) + with patch_ingress(): + app = build_openai_app({"llm_configs": [llm_config]}) + handle = serve.run(app, name="kv_offload_gpu_test") + yield handle + serve.shutdown() + + async def _discover_replicas(self, handle): + endpoints = {} + for _ in range(100): + async with handle.choose_replica() as selection: + replica = selection._replica + if replica.backend_http_endpoint is not None: + endpoints[ + replica.replica_id.to_full_id_str() + ] = replica.backend_http_endpoint + if len(endpoints) == NUM_REPLICAS: + return endpoints + await asyncio.sleep(0.5) + raise AssertionError( + f"Expected {NUM_REPLICAS} replicas with backend endpoints, " + f"found {len(endpoints)}." + ) + + @pytest.mark.asyncio + @pytest.mark.timeout(600) + async def test_offload_routes_to_cpu_prefix_and_reloads(self, deployed_handle): + """A GPU-evicted CPU prefix stays routable and reloads on its replica.""" + router = serve.get_deployment_handle( + "LLMRouter", app_name="kv_offload_gpu_test" + ) + replica_endpoints = await self._discover_replicas(deployed_handle) + + async def all_registered(): + registered = await router.get_kv_event_worker_replicas.remote() + return sorted(registered.values()) == sorted(replica_endpoints) + + await async_wait_for_condition(all_registered, timeout=90) + replica_by_worker = await router.get_kv_event_worker_replicas.remote() + endpoints = { + worker_id: replica_endpoints[replica_id] + for worker_id, replica_id in replica_by_worker.items() + } + cached_worker, miss_worker = sorted(endpoints) + target_tokens = tokenize_prompt(endpoints[cached_worker], OFFLOAD_MESSAGES) + target_blocks = num_prompt_blocks(target_tokens) + assert 4 < target_blocks < 32 + + post_chat(endpoints[cached_worker], OFFLOAD_MESSAGES, max_tokens=2) + + async def target_is_on_gpu(): + scores = await router.get_kv_overlap_scores.remote(target_tokens) + return scores.get(cached_worker, {}).get("device_blocks") == target_blocks + + await async_wait_for_condition(target_is_on_gpu, timeout=60) + + # Each unique prompt displaces old GPU blocks. Native offload retains + # the target in CPU memory and emits CPU-tier events as that happens. + async def target_is_only_on_cpu(): + scores = await router.get_kv_overlap_scores.remote(target_tokens) + cached_score = scores.get(cached_worker, {}) + return ( + cached_score.get("device_blocks") == 0 + and cached_score.get("host_pinned_blocks", 0) > 0 + ) + + offloaded = False + for i in range(12): + filler = [ + { + "role": "user", + "content": ( + f"Unique GPU eviction sequence {i}. " + f"filler-{i} " * 60 + ), + } + ] + post_chat(endpoints[cached_worker], filler, max_tokens=1) + for _ in range(10): + if await target_is_only_on_cpu(): + offloaded = True + break + await asyncio.sleep(0.5) + if offloaded: + break + if not offloaded: + raise AssertionError("Target prefix was not offloaded from GPU to CPU.") + + scores = await router.get_kv_overlap_scores.remote(target_tokens) + cached_score = scores[cached_worker] + miss_score = scores[miss_worker] + assert cached_score["host_pinned_extension_blocks"] > 0 + assert ( + 0 + < cached_score["router_credit_blocks"] + < cached_score["host_pinned_blocks"] + ) + assert miss_score["device_blocks"] == 0 + assert miss_score["host_pinned_blocks"] == 0 + + # Exercise the production HAProxy -> LLMRouter -> KVAwareRouter path. + # The response can only report a cache hit if it reached the replica + # whose target prefix is still available in CPU memory. + response = post_chat(("127.0.0.1", 8000), OFFLOAD_MESSAGES, max_tokens=2) + assert response["usage"]["prompt_tokens_details"]["cached_tokens"] > 0 + + # The request was served from CPU and the loaded blocks are visible on + # GPU again on that replica, while the uncached replica remains empty. + await async_wait_for_condition(target_is_on_gpu, timeout=60) + scores = await router.get_kv_overlap_scores.remote(target_tokens) + assert scores[miss_worker]["device_blocks"] == 0 + assert scores[miss_worker]["host_pinned_blocks"] == 0 + + class TestKvScoring: """End-to-end KV-aware routing: a request routed through a deployed KVAwareRouter is scored by the selection service and lands on a live diff --git a/release/llm_tests/kv_router_test/utils.py b/release/llm_tests/kv_router_test/utils.py index a2297aaec629..64c93f803459 100644 --- a/release/llm_tests/kv_router_test/utils.py +++ b/release/llm_tests/kv_router_test/utils.py @@ -144,6 +144,13 @@ def get_registered_worker_ids(self): async def get_kv_overlap_blocks(self, token_ids): """(Test only) Per-worker device-tier KV overlap blocks for a sequence.""" + 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): + """(Test only) Per-worker overlap across every KV storage tier.""" svc = self._kv_token_tracker._svc if svc is None: return {} @@ -154,7 +161,7 @@ async def get_kv_overlap_blocks(self, token_ids): "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_worker_active_requests(self, worker_id): """(Test only) In-flight requests the service tracks as active load on From f41c117c54b1b31cfb65560735fb53b8e0dd4239 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Mon, 27 Jul 2026 17:40:27 -0700 Subject: [PATCH 2/4] Self review Signed-off-by: Jeffrey Wang --- .../deployments/routers/kv/test_kv_events.py | 32 ++++------ .../kv_router_test/test_kv_events.py | 61 ++++++------------- release/llm_tests/kv_router_test/utils.py | 20 ++++++ 3 files changed, 50 insertions(+), 63 deletions(-) diff --git a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py index 4d12979be357..c4521a3acace 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py +++ b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py @@ -64,29 +64,19 @@ def test_configure_enables_events_and_pins_runtime_env(self): assert llm_config.runtime_env["env_vars"]["VLLM_USE_SIMPLE_KV_OFFLOAD"] == "0" assert llm_config.runtime_env["env_vars"]["EXISTING_ENV"] == "value" - def test_native_offload_enables_self_describing_events(self): - """The resolved native connector emits complete CPU-tier KV events.""" - extra_config = {"offload_prompt_only": False} - vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - kv_offloading_size=2.0, - kv_offloading_backend="native", - ), - kv_transfer_config=SimpleNamespace(kv_connector_extra_config=extra_config), - ) - - enable_native_kv_offload_events(vllm_config) - - assert extra_config == { - "offload_prompt_only": False, - "self_describing_kv_events": True, - } - - def test_no_offload_is_unchanged(self): + @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=None, + kv_offloading_size=offload_size, kv_offloading_backend="native", ), kv_transfer_config=SimpleNamespace(kv_connector_extra_config=extra_config), @@ -94,7 +84,7 @@ def test_no_offload_is_unchanged(self): enable_native_kv_offload_events(vllm_config) - assert extra_config == {"existing": "value"} + assert extra_config == expected @pytest.mark.parametrize( "engine_kwargs, local_rank, expected_port, expected_replay_port", diff --git a/release/llm_tests/kv_router_test/test_kv_events.py b/release/llm_tests/kv_router_test/test_kv_events.py index f83d51ac4756..b28675e0c830 100644 --- a/release/llm_tests/kv_router_test/test_kv_events.py +++ b/release/llm_tests/kv_router_test/test_kv_events.py @@ -10,10 +10,15 @@ from ray import serve from ray._common.test_utils import async_wait_for_condition from ray.serve.config import RequestRouterConfig -from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_openai_app +from ray.serve.llm import LLMConfig, ModelLoadingConfig from ray.serve.llm.request_router import KVAwareRouter -from utils import _TestKVAwareRouter, patch_ingress +from utils import ( + _TestKVAwareRouter, + build_kv_app, + discover_replica_endpoints, + patch_ingress, +) MODEL_ID = "qwen3-0.6b" MODEL_SOURCE = "Qwen/Qwen3-0.6B" @@ -151,28 +156,11 @@ def deployed_handle(self): # Swap the ingress for the introspection LLMRouter so the embedded # tracker's state is reachable over the deployment handle. with patch_ingress(): - app = build_openai_app({"llm_configs": [llm_config]}) + app = build_kv_app(llm_config) handle = serve.run(app, name=APP_NAME) yield handle serve.shutdown() - async def _discover_replicas(self, handle): - """Map each replica's full id to its backend HTTP endpoint.""" - endpoints = {} - for _ in range(100): - async with handle.choose_replica() as selection: - replica = selection._replica - if replica.backend_http_endpoint is not None: - replica_id = replica.replica_id.to_full_id_str() - endpoints[replica_id] = replica.backend_http_endpoint - if len(endpoints) == NUM_REPLICAS: - return endpoints - await asyncio.sleep(0.5) - raise AssertionError( - f"Expected {NUM_REPLICAS} replicas with backend endpoints, " - f"found {len(endpoints)}." - ) - @pytest.mark.asyncio @pytest.mark.timeout(600) async def test_kv_events_reach_selection_service(self, deployed_handle): @@ -182,7 +170,9 @@ async def test_kv_events_reach_selection_service(self, deployed_handle): worker.""" router = serve.get_deployment_handle("LLMRouter", app_name=APP_NAME) - replica_endpoints = await self._discover_replicas(deployed_handle) + replica_endpoints = await discover_replica_endpoints( + deployed_handle, NUM_REPLICAS + ) # Each replica advertises its KV-events endpoint via record_routing_stats; # the controller propagates it on the LongPoll replica snapshot and the @@ -268,7 +258,9 @@ async def test_chat_tokens_match_prefill(self, deployed_handle): """Ensure chat template is applied: a chat request scores the same overlap as the prompt rendered with the model's chat template and tokenized as raw text.""" router = serve.get_deployment_handle("LLMRouter", app_name=APP_NAME) - replica_endpoints = await self._discover_replicas(deployed_handle) + replica_endpoints = await discover_replica_endpoints( + deployed_handle, NUM_REPLICAS + ) async def all_registered(): registered = await router.get_kv_event_worker_replicas.remote() @@ -349,28 +341,11 @@ def deployed_handle(self): log_engine_metrics=False, ) with patch_ingress(): - app = build_openai_app({"llm_configs": [llm_config]}) + app = build_kv_app(llm_config) handle = serve.run(app, name="kv_offload_gpu_test") yield handle serve.shutdown() - async def _discover_replicas(self, handle): - endpoints = {} - for _ in range(100): - async with handle.choose_replica() as selection: - replica = selection._replica - if replica.backend_http_endpoint is not None: - endpoints[ - replica.replica_id.to_full_id_str() - ] = replica.backend_http_endpoint - if len(endpoints) == NUM_REPLICAS: - return endpoints - await asyncio.sleep(0.5) - raise AssertionError( - f"Expected {NUM_REPLICAS} replicas with backend endpoints, " - f"found {len(endpoints)}." - ) - @pytest.mark.asyncio @pytest.mark.timeout(600) async def test_offload_routes_to_cpu_prefix_and_reloads(self, deployed_handle): @@ -378,7 +353,9 @@ async def test_offload_routes_to_cpu_prefix_and_reloads(self, deployed_handle): router = serve.get_deployment_handle( "LLMRouter", app_name="kv_offload_gpu_test" ) - replica_endpoints = await self._discover_replicas(deployed_handle) + replica_endpoints = await discover_replica_endpoints( + deployed_handle, NUM_REPLICAS + ) async def all_registered(): registered = await router.get_kv_event_worker_replicas.remote() @@ -499,7 +476,7 @@ def kv_aware_handle(self): ), log_engine_metrics=False, ) - app = build_openai_app({"llm_configs": [llm_config]}) + app = build_kv_app(llm_config) handle = serve.run(app, name="kv_scoring_gpu_test") yield handle serve.shutdown() diff --git a/release/llm_tests/kv_router_test/utils.py b/release/llm_tests/kv_router_test/utils.py index 64c93f803459..6f2e9ce8c82d 100644 --- a/release/llm_tests/kv_router_test/utils.py +++ b/release/llm_tests/kv_router_test/utils.py @@ -7,6 +7,7 @@ exposes the tracker's state as handle-callable methods. """ +import asyncio from contextlib import contextmanager from dataclasses import asdict import sys @@ -69,6 +70,25 @@ def build_kv_app(llm_config): return build_openai_app({"llm_configs": [llm_config]}) +async def discover_replica_endpoints(handle, expected_replicas): + """Map each replica id to its direct-ingress HTTP endpoint.""" + endpoints = {} + for _ in range(100): + async with handle.choose_replica() as selection: + replica = selection._replica + if replica.backend_http_endpoint is not None: + endpoints[ + replica.replica_id.to_full_id_str() + ] = replica.backend_http_endpoint + if len(endpoints) == expected_replicas: + return endpoints + await asyncio.sleep(0.5) + raise AssertionError( + f"Expected {expected_replicas} replicas with backend endpoints, " + f"found {len(endpoints)}." + ) + + class _TestKVAwareRouter(RoundRobinRouter, KVAwareRouter): """A ``KVAwareRouter`` subclass that borrows ``RoundRobinRouter``'s selection. From 05b205fe04134993567785982f5e3d6e02dd27fc Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Wed, 29 Jul 2026 10:08:16 -0700 Subject: [PATCH 3/4] Allow KV offloading for any routers Signed-off-by: Jeffrey Wang --- .../deployments/routers/kv/test_kv_events.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py index c4521a3acace..812250cb65c9 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py +++ b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py @@ -13,6 +13,9 @@ get_kv_event_routing_stats, resolve_kv_event_source_endpoint, ) +from ray.llm._internal.serve.routing_policies.kv_aware.utils import ( + _maybe_setup_kv_aware_routing, +) from ray.serve.llm.request_router import KVAwareRouter @@ -64,6 +67,31 @@ def test_configure_enables_events_and_pins_runtime_env(self): assert llm_config.runtime_env["env_vars"]["VLLM_USE_SIMPLE_KV_OFFLOAD"] == "0" assert llm_config.runtime_env["env_vars"]["EXISTING_ENV"] == "value" + def test_non_kv_aware_router_preserves_kv_offload_config(self): + """Non-KV-aware routers retain the user's KV offload configuration.""" + llm_config = make_kv_aware_llm_config( + runtime_env={ + "env_vars": { + "EXISTING_ENV": "value", + "VLLM_USE_SIMPLE_KV_OFFLOAD": "1", + } + }, + engine_kwargs={ + "kv_offloading_size": 2.0, + "kv_offloading_backend": "native", + }, + ) + llm_config.deployment_config["request_router_config"][ + "request_router_class" + ] = "ray.serve.experimental.consistent_hash_router:ConsistentHashRouter" + + _maybe_setup_kv_aware_routing(llm_config.deployment_config, llm_config) + + assert "kv_events_config" not in llm_config.engine_kwargs + assert llm_config.engine_kwargs["kv_offloading_size"] == 2.0 + assert llm_config.engine_kwargs["kv_offloading_backend"] == "native" + assert llm_config.runtime_env["env_vars"]["VLLM_USE_SIMPLE_KV_OFFLOAD"] == "1" + @pytest.mark.parametrize( "offload_size, expected", [ From 693d0b71399ae1f7eddb0a42ad0d9662a1d096ea Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Fri, 31 Jul 2026 01:44:23 -0700 Subject: [PATCH 4/4] [serve][llm] Fix KV offload lint Signed-off-by: Jeffrey Wang --- .../serve/cpu/deployments/routers/kv/test_kv_events.py | 6 +++--- release/llm_tests/kv_router_test/test_kv_events.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py index 812250cb65c9..9424ab0996f8 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py +++ b/python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_events.py @@ -6,6 +6,9 @@ import ray from ray.llm._internal.serve.core.configs.llm_config import LLMConfig +from ray.llm._internal.serve.routing_policies.kv_aware.utils import ( + _maybe_setup_kv_aware_routing, +) from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import ( assign_replica_kv_events_endpoint, configure_kv_events_for_kv_routing, @@ -13,9 +16,6 @@ get_kv_event_routing_stats, resolve_kv_event_source_endpoint, ) -from ray.llm._internal.serve.routing_policies.kv_aware.utils import ( - _maybe_setup_kv_aware_routing, -) from ray.serve.llm.request_router import KVAwareRouter diff --git a/release/llm_tests/kv_router_test/test_kv_events.py b/release/llm_tests/kv_router_test/test_kv_events.py index b28675e0c830..af262b68a7bf 100644 --- a/release/llm_tests/kv_router_test/test_kv_events.py +++ b/release/llm_tests/kv_router_test/test_kv_events.py @@ -10,7 +10,7 @@ from ray import serve from ray._common.test_utils import async_wait_for_condition from ray.serve.config import RequestRouterConfig -from ray.serve.llm import LLMConfig, ModelLoadingConfig +from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_openai_app from ray.serve.llm.request_router import KVAwareRouter from utils import (