[jeffreywang88]

Description

Add session pinning to KVAwareRouter so conversation turns reuse the same replica and its KV cache.

The first request selects a replica using estimated token load and KV-cache overlap. Later requests with the same x-session-id stay on that replica, even when another has lower load or greater overlap. If the replica becomes ineligible or the session expires, the router selects again.

Session affinity flow

1. First turn reaches ingress A. Ray calls its local selection service's select_and_reserve(...) with the session ID and allowed replicas. With no binding yet, the service selects by token load and KV overlap, books the request's load, and binds the session to that replica.
2. A broadcasts load and binding separately. The reservation goes out as a load_reservation update. The binding comes from the session_binding_updates() stream, which emits each new binding and each TTL refresh, and goes out as a session_affinity update.
3. Ingress B applies both.
   • apply_session_binding(...) reconciles the binding by version.
   • create_reservation(...) mirrors the load on the original worker, unless the request already completed. A completion that arrives mid-booking frees the booking.
   • B always learns the binding, even when it skips the load.
4. A later turn reaches B. B follows the replicated session binding even if another worker has lower load or more overlap. If the bound worker is unavailable or excluded from the allowlist, it rebinds the session with a newer version, and every peer receives it through the stream. Each turn books and broadcasts its own load, and each completed turn refreshes the session's TTL on every peer.

API shape

Session pinning defaults to a 30-minute idle timeout. Configure session_affinity_ttl_s through request_router_kwargs, or set it to 0 to disable pinning.

from ray.serve.llm import LLMConfig
from ray.serve.llm.request_router import KVAwareRouter

llm_config = LLMConfig(
    model_loading_config={
        "model_id": "qwen3-0.6b",
        "model_source": "Qwen/Qwen3-0.6B",
    },
    deployment_config={
        "request_router_config": {
            "request_router_class": KVAwareRouter,
            "request_router_kwargs": {
                "session_affinity_ttl_s": 600,  # 10-minute idle timeout
            },
        },
    },
)

Benchmark

For long-context, long-horizon workloads like AgentX, routing a request to a different replica later in a session can trigger substantial redundant prefill, significantly hurting throughput and interactivity.

b6f2a69:

Related issues

Dynamo PR: ai-dynamo/dynamo#16969.

[gemini-code-assist]

Code Review

This pull request introduces session pinning to the KVAwareRouter and KVTokenTracker using an x-session-id header to improve KV cache reuse, along with a configurable idle timeout session_affinity_ttl_s. Feedback on the changes highlights potential runtime issues, including an AttributeError if request_router_kwargs is None in build_kv_token_tracker, and a TypeError if session_affinity_ttl_s is passed as None in KVAwareRouter.initialize_state. Additionally, it is recommended to validate that session_affinity_ttl_s is non-negative within the KVTokenTracker constructor.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 40fc1b3. Configure here.}

[jeffreywang88]

Note for reviewers

This is only a rename to match dynamo's current API (it renamed tenancy to routing groups). Behavior doesn't change.

In dynamo, routing_group and model_name together pick a routing partition. Each partition has its own workers, KV index, load tracking, and session-affinity table. Ray doesn't use routing groups: every call passes "default", so each SelectionService has one partition.

LoRA doesn't use routing groups either. Serve multiplexing narrows the candidate replicas to those that have the adapter, and each request's lora_name salts its KV block hashes. That keeps each adapter's cache separate inside the same partition.

[jthomson04]

I benchmarked this path with ai-dynamo/dynamo#16969 (details on that PR: 4 ingress processes, 1,024 mock vLLM workers, AgentX sessions, 1,000 ongoing selections per ingress).

_apply_replica_updates awaits one create_reservation / apply_session_binding per update, so it advances about one update per event-loop iteration. With the ingress loaded:

• each replica applied about 240 updates/s against about 730/s arriving;
• lag grew to 122 s;
• every peer reservation was skipped as already completed, so each replica's load view held only its own requests.

The cache hit ratio at 4 replicas was 0.64, against 0.93 for one replica, and KV churn was 11× higher. Awaiting each round together (asyncio.gather over the round's updates) brought it to 0.89. With the binding-stream batching on the Dynamo side, it reached 0.925.

Without session ids the same backlog appears (13–18 s of lag, 72% of reservations skipped), so this is independent of session pinning. _apply_reservations on master has the same one-at-a-time loop.

Minor: choose_replicas recomputes blake2b for every candidate on every request, about 0.5–0.6 ms of event-loop time per request with 1,024 replicas. _replica_id_by_worker is already kept current from LongPoll, so the inverse map could be maintained there.

[jeffreywang88]

Thanks @jthomson04! I was able to repro the issue by using ray serve LLM with 4 ingress replicas, 64 mocker engines (LLMServer running MockerEngine, each with a in-process dynamo mocker), and 640 concurrent sessions replaying AgentX at 850 requests/s and confirmed that b6f2a69 and ai-dynamo/dynamo@fcafcc7 together resolves the issue.

asyncio.gather over a round's update + _replica_id_by_worker cache are implemented in b6f2a69. Will add the O(1) inverse map in a separate follow-up. Also added test_session_bindings_forwarded_without_drops to verify that we don't drop session binding updates.

[jeffreywang88]

Note for reviewers: This validate the SOFT affinity mode -- when the allowlist excludes the original bound worker, selection service rebinds the session to another eligible worker.

[jeffreywang88]

Need to build a wheel with the upstream changes for KV release tests.