From cef80d7219eb25fb3c4763a5ae87c26aaed11699 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Sun, 2 Aug 2026 13:33:23 -0700 Subject: [PATCH] [llm][kv][13/N] Enable KV cache offloading & KV routing awareness of CPU KV caches (#65063) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description This PR adds native vLLM KV-cache offloading support to ray serve LLM KV-aware routing. KV-aware deployments use `OffloadingConnector` (enabled via `VLLM_USE_SIMPLE_KV_OFFLOAD=0` and `kv_connector_extra_config["self_describing_kv_events"] = True`) to publish CPU-tier KV events, letting the selection service account for CPU-resident prefixes and route reloads to the appropriate replica. ### KV event flow `OffloadingConnector` → `BlockStored` / `BlockRemoved` (medium="CPU") ([reference](https://github.com/vllm-project/vllm/blob/9d1aa4dda85a58c74b23ff79aa4a20594ee6f359/vllm/distributed/kv_events.py#L61)) → vLLM `KVEventBatch` over ZMQ PUB → Dynamo listener decodes and normalizes events → CPU events indexed as HostPinned tier → selection service scores GPU + CPU prefix overlap per replica ### API shape ``` from ray.serve.config import RequestRouterConfig from ray.serve.llm import LLMConfig from ray.serve.llm.request_router import KVAwareRouter llm_config = LLMConfig( model_loading_config={ "model_id": "my-model", "model_source": "my-model", }, deployment_config={ "request_router_config": RequestRouterConfig( request_router_class=KVAwareRouter, ), }, engine_kwargs={ "enable_prefix_caching": True, "kv_offloading_backend": "native", "kv_offloading_size": 8, # GiB of CPU cache per replica }, ) ``` ## Validation - **CI tests**: Added unit and release coverage for CPU-tier event ingestion, scoring, and end-to-end offload/reload behavior. Non-KV-aware routers retain user-specified KV offload configuration. - Observed KV cache offloading / reloading, and I'll add grafana dashboards for KV cache offloading observability in a follow-up. - Workload characteristic: A two-replica Ray Serve LLM workload on L4s using Qwen3-0.6B with native KV offload, capped GPU KV cache, sticky session routing, long repeated prompts, forced prefix-cache resets, and same-session follow-ups. Screenshot 2026-07-28 at 11 33 43 PM ## Related issues > Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- 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 | 66 ++++++- .../kv_router_test/test_kv_event_ingestion.py | 109 ++++++++++- .../kv_router_test/test_kv_events.py | 185 +++++++++++++++--- release/llm_tests/kv_router_test/utils.py | 29 ++- 6 files changed, 382 insertions(+), 34 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 190dbe05cbdb..2e0fd3172e6b 100644 --- a/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py +++ b/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py @@ -63,6 +63,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, get_prompt_token_routing_stats, get_token_channel_endpoints, @@ -589,6 +590,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 f4f83ffb7584..49c433f058de 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 @@ -42,20 +42,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 9eebc8ea4543..33a7dd5b78d5 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,9 +6,13 @@ 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, + enable_native_kv_offload_events, get_kv_event_routing_stats, get_prompt_token_routing_stats, get_token_channel_endpoints, @@ -43,9 +47,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"] == { @@ -55,6 +66,55 @@ 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_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", + [ + (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", 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 5c865a66553d..501f1bc4609e 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 @@ -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 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..af262b68a7bf 100644 --- a/release/llm_tests/kv_router_test/test_kv_events.py +++ b/release/llm_tests/kv_router_test/test_kv_events.py @@ -13,7 +13,12 @@ from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_openai_app 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" @@ -44,6 +49,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): @@ -141,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): @@ -172,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 @@ -258,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() @@ -296,6 +298,145 @@ 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_kv_app(llm_config) + handle = serve.run(app, name="kv_offload_gpu_test") + yield handle + serve.shutdown() + + @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 discover_replica_endpoints( + deployed_handle, NUM_REPLICAS + ) + + 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 @@ -335,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 487a46a1a3f2..d31135e169f3 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 @@ -79,6 +80,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. @@ -174,6 +194,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 {} @@ -184,7 +211,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