[Router][Feat] loadaware routing: KV-cache-aware placement weighted by live load

[bazakeliad]

Motivation

kvaware maximizes cache-hit benefit alone: it routes to the first instance reported to hold a request's prefix, however busy that instance is. Under workloads with popular shared prefixes this concentrates load — every request for a hot prefix lands on the one engine holding it, which queues while its peers idle. Cache affinity and load balance pull in opposite directions, and the stock router only pulls one way.

What this PR adds

A loadaware routing logic that scores every endpoint

score(i) = matched_tokens(i) / prompt_tokens - beta * relative_load(i)
relative_load(i) = (load(i) - mean_load) / max(1, mean_load)

and routes to the argmax, so a warm-but-saturated instance can lose to a cold-but-idle one.

Design properties:

• Both terms are dimensionless — a fraction of this prompt, and a signed fraction of this fleet's mean in-flight load, recomputed per request from the router's own request_stats. The single tunable beta is therefore a pure exchange rate between cache locality and load balance and carries no unit from the deployment, which is what lets it ship a meaningful default (1.0: an endpoint 100% above fleet-average load is docked one full cache hit's worth of preference).
• KvawareRouter is untouched. LoadAwareRouter subclasses it and overrides only the selection step. The rest is a new enum member, a factory branch, --loadaware-beta (with LOADAWARE_BETA env fallback), and the parser choice.
• Works against stock LMCache lookups. It scores whatever holders layout_info reports; a multi-holder lookup on the LMCache side (their acknowledged lookup() TODO) widens the benefit term to the whole fleet, and I plan to submit that separately to LMCache.
• No new dependencies, no new stats-collection path (request_stats was already handed to every route_request call).

Measured results

2×NVIDIA A10 OpenShift cluster serving Qwen/Qwen2.5-3B-Instruct through this stack (chart 0.1.11), Zipf-distributed shared prefixes (128 × 2048 tokens, s = 0.9), 16 req/s open-loop Poisson, 500 requests × 20 paired seeds per arm, exact Wilcoxon signed-rank on per-seed differences:

Metric	kvaware	loadaware β=0.5	Change
Load imbalance (busiest/idlest vllm:num_requests_running)	2.358	1.249	−48.1%, p < 0.0001
Same, independent re-run two days later	2.452	1.272	−49.4%, p < 0.0001
β = 0 ablation (imbalance)	2.358	2.662	null, p = 0.97 — the load term is the entire mechanism
vLLM prefix-cache hit rate	91.2%	90.7%	−0.5 pp locality cost
Requests missing a 150 ms TTFT objective	—	—	−19.0%, p = 0.0021
KV-cache memory spread across engines	1.70×	1.18×	the cache itself balances, not just request counts
GPU util / power / router CPU / throughput	—	—	unchanged

Client-observed TTFT p95 was a pre-registered co-primary and did not reach significance at this operating point (−2.7%, p = 0.12): the fleet never queued (vllm:num_requests_waiting = 0 in all scrapes), so there was no queueing delay for placement to remove. The imbalance reduction is the claim; the latency null is its boundary. Full methodology, per-seed data, and reproduction scripts: https://github.com/BenEpstein/caching-in-llms.

Tests

28 offline unit tests in src/tests/test_loadaware_router.py covering the score arithmetic (scale invariance, fleet-relative load, benefit normalization and cap, the near-idle clamp), the instance-id → URL bridge (engine-restart staleness, phantom credit from dead instance ids), tie-breaking determinism, fallback behavior, and the beta configuration chain (flag > env > default). They run without lmcache installed.

---

• Make sure the code changes pass the pre-commit checks.
• Sign-off your commit by using -s when doing git commit
• Try to classify PRs for easy understanding of the type of changes, such as [Bugfix], [Feat], and [CI].

[gemini-code-assist]

Code Review

This pull request introduces a new loadaware routing logic (LoadAwareRouter) that balances cache-hit benefits against relative engine loads using a tunable beta parameter. The feedback focuses on performance and robustness improvements in the new router: making tokenize_prompt asynchronous to avoid blocking the event loop with synchronous HTTP requests, optimizing refresh_instance_map to query endpoints concurrently and parse IPs more robustly, and adding a guard in route_request to handle empty endpoints gracefully.

[ruizhang0101]

This is great! Could you make the docstring concise? Also could you add a user-facing documentation(i.e. rst files) for this as well?

[bazakeliad]

Thanks for the review, @ruizhang0101! Both points addressed in bbb610e:

• Trimmed the LoadAwareRouter class and method docstrings to the essentials.
• Added a user-facing tutorial at docs/source/use_cases/loadaware-routing.rst (linked from the Use Cases toctree), covering deployment, how the scoring works, and tuning beta — the design rationale that used to live in the docstring moved there.

[ruizhang0101]

LGTM