[serve][llm] Session pinning for KVAwareRouter

doc/source/serve/llm/user-guides/kv-aware-routing.md

@@ -26,7 +26,7 @@ The best policy depends on your workload. You can configure different routers th
 | `RoundRobinRouter` | Request order | Prompts share little beyond the system prompt, and you want the simplest even spread. | No KV cache awareness. |
 | `ConsistentHashRouter` | Hash of the `x-session-id` header | Clients can provide a session ID and each session’s history forms the reusable prefix. | A session with more turns or longer input or output sequences can saturate its replica. |
 | `PrefixCacheAffinityRouter` | Prompt text matched against a router-maintained prefix tree | Requests frequently share a long textual prefixes. | Approximates engine KV cache state from prompt text and falls back to power of two choices when queue lengths become imbalanced. |
-| `KVAwareRouter` | Token load: the request's remaining prefill work after accounting for KV cache overlap, plus the replica's active prefill and decode work | Replicas carry uneven token load, GPU memory is under pressure, or requests share prefixes beyond the system prompt. | Requires direct streaming, and its extra overhead may not pay off on simple, uniform workloads. |
+| `KVAwareRouter` | Token load: the request's remaining prefill work after accounting for KV cache overlap, plus the replica's active prefill and decode work, with session pinning through `x-session-id` | Replicas carry uneven token load, GPU memory is under pressure, or requests share prefixes beyond the system prompt. | Requires direct streaming, and its extra overhead may not pay off on simple, uniform workloads. |
 
 ## Installation
 
@@ -101,6 +101,30 @@ Run `serve run config.yaml`.
 :::
 ::::
 
+### Use session pinning
+
+Send the same `x-session-id` header on each turn of a conversation to improve KV cache reuse:
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -H "x-session-id: chat-123" \
+  -d '{"model": "qwen3-0.6b", "messages": [{"role": "user", "content": "Hello!"}]}'
+```
+
+The first request selects a replica by estimated token load and KV cache overlap. Later turns stay on that replica even when another has lower load or more cached tokens. After the session's idle timeout, or if the replica becomes unavailable or ineligible, the router selects again.
+
+Session pinning is enabled by default with a 30-minute idle timeout. Set `session_affinity_ttl_s` in seconds to change it:
+
+```python
+request_router_config = RequestRouterConfig(
+    request_router_class=KVAwareRouter,
+    request_router_kwargs={"session_affinity_ttl_s": 600},
+)
+```
+
+Set it to `0` to disable pinning. Requests without `x-session-id` use token-load scoring and KV cache overlap.
+
 ### Tuning
 
 Set these in the cluster environment:

python/ray/llm/_internal/serve/core/ingress/router.py

@@ -163,7 +163,7 @@ async def __init__(
             self._kv_token_tracker = build_kv_token_tracker(
                 llm_config, server.deployment_id
             )
-            self._kv_token_tracker.start_reservation_broadcast(get_llm_router_handle())
+            self._kv_token_tracker.start_replica_broadcast(get_llm_router_handle())
             # Lazy import: this module pulls in vLLM's renderer;
             # keep it off the non-KV ingress import path.
             from ray.llm._internal.serve.routing_policies.kv_aware.vllm.tokenizer import (  # noqa: E501
@@ -300,9 +300,9 @@ async def on_lifecycle_events(self, batch):
         """
         return await self._kv_token_tracker.on_lifecycle_events(batch)
 
-    async def on_reservations_created(self, batch):
-        """Ingress-facing intake for already-selected reservation bookings."""
-        return await self._kv_token_tracker.on_reservations_created(batch)
+    async def on_replica_updates(self, source_writer_id, batch):
+        """Accept load and affinity updates from peer ingress replicas."""
+        return await self._kv_token_tracker.on_replica_updates(source_writer_id, batch)
 
     def _push_prompt_tokens(
         self,

python/ray/llm/_internal/serve/routing_policies/kv_aware/constants.py

@@ -56,6 +56,13 @@
 # transient actor outage) would otherwise leave its entry tracked forever.
 REQUEST_TRACKING_TTL_S = 3600
 
+# Session pinning idle timeout in seconds.
+DEFAULT_SESSION_AFFINITY_TTL_S = 1800.0
+
+# Soft affinity rebinds a session when Serve excludes its bound worker from the
+# request's candidates; hard affinity would reject the request instead.
+SESSION_AFFINITY_MODE = "soft"
+
 # Bound best-effort lifecycle broadcast delivery so one slow ingress replica
 # cannot stall an engine replica's whole lifecycle-event queue indefinitely.
 LIFECYCLE_EVENT_BROADCAST_TIMEOUT_S = 3

python/ray/llm/_internal/serve/routing_policies/kv_aware/kv_aware_router.py

@@ -4,6 +4,7 @@
 
 from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
 from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
+    DEFAULT_SESSION_AFFINITY_TTL_S,
     REQUEST_TOKEN_IDS_KWARG,
 )
 from ray.llm._internal.serve.routing_policies.kv_aware.kv_token_tracker import (
@@ -42,12 +43,25 @@ class KVAwareRouter(MultiplexMixin, RequestRouter):
     Dynamo selection service and the global KV index) built by the LLMRouter in
     this same ingress process; this per-handle router stays thin and simply maps
     candidate replicas to/from Dynamo worker ids.
+
+    Send the same ``x-session-id`` header on each conversation turn to keep
+    requests on the same replica. The first turn uses token-load scoring and
+    KV-cache overlap. Later turns stay on that replica until the session's idle
+    timeout expires or the replica becomes unavailable or ineligible.
     """
 
-    def initialize_state(self):
+    def initialize_state(
+        self,
+        session_affinity_ttl_s: float = DEFAULT_SESSION_AFFINITY_TTL_S,
+    ):
         """Bind to the ``KVTokenTracker`` the LLMRouter registered in this
         process. When absent (e.g. the proxy's fallback router), KV-aware
         routing degrades to load-balanced selection instead of erroring.
+
+        Args:
+            session_affinity_ttl_s: How long to keep an idle session on the same
+                replica, in seconds. Defaults to 1800 (30 minutes).
+                ``0`` disables session pinning.
         """
         self._kv_token_tracker = get_kv_token_tracker()
         if self._kv_token_tracker is None:
@@ -124,6 +138,7 @@ async def choose_replicas(
             # the selection service salts its KV hashes with it, so each adapter
             # gets its own KV namespace. Empty means the base model.
             lora_name=pending_request.metadata.multiplexed_model_id or None,
+            session_id=pending_request.metadata.session_id or None,
         )
         return [[worker_id_to_replica[selection["worker_id"]]]]