feat: cache aware per request affinity by sdevare-nv · Pull Request #238 · vllm-project/router

sdevare-nv · 2026-09-04T00:43:39Z

Includes #216 unmodified, and fixes a second, opposite-direction failure that #216
exposes.

The problem

CacheAwarePolicy::select_worker_with_headers decided whether to use prefix
affinity by looking at the whole fleet:

let is_imbalanced = max_load.saturating_sub(min_load) > balance_abs_threshold
    && (max_load as f32) > (min_load as f32 * balance_rel_threshold);
if is_imbalanced {
    return self.select_worker_min_load(...);   // returns before the tree is consulted
}

The question it asks is "is any worker hot?", and if the answer is yes it discards
cache affinity for every request — including requests whose preferred worker is
completely idle. One runaway worker turns off prefix routing for the entire fleet.

Under prefill/decode disaggregation that gate is open almost permanently. Prefill
worker load counts queued requests as well as running ones, so at high concurrency
the spread between the busiest and least busy prefill worker clears the threshold
continuously. Routing degenerates to shortest-queue, the prefix tree stops being
consulted, and prompts that were already cached get recomputed from scratch. That
saturates prefill, which in turn starves decode.

This is latent on main today. While the health checker zeroes load counters every
ten cycles — the bug #216 fixes — max_load - min_load stays near zero, the gate
never fires, and routing is pure affinity by accident. Making the counters honest
turns the gate on for the first time, so #216 on its own is a significant
regression for prefill-heavy workloads. That is why the two changes are proposed
together rather than separately.

The fix

Evaluate the same thresholds per request, against the worker that request
actually wants, rather than against the fleet maximum:

Some(idx) if self.is_worker_overloaded(workers[idx].load(), min_load) => {
    // the worker holding this prefix is the runaway: pay the miss, shed the load
    RouterMetrics::record_load_balancing_event();
    healthy_indices.iter().min_by_key(|&&i| workers[i].load()).copied()
}
other => other,   // otherwise keep affinity

Shedding load only helps when the worker holding the prefix is itself the hot one,
so that is now the only case where affinity is given up. A worker that is merely
busier than its peers keeps serving its cached prefixes; a genuine hot spot still
gets steered away from, exactly as before. select_worker_min_load has no
remaining callers and is removed.

Three regression tests come with it. The load-bearing one is
test_affinity_survives_unrelated_hot_worker: a runaway worker elsewhere in the
fleet must not cost affinity for a request whose own worker is idle — the case the
old fleet-wide gate got wrong. Its complement,
test_affinity_yields_when_cached_worker_is_the_hot_one, pins the case where the
cache miss should be paid, so the fix cannot silently disable load shedding.

Validated end to end on a prefill-heavy agentic benchmark under P/D
disaggregation: with this change the prefill prefix cache hit rate and overall
throughput return to their pre-#216 behaviour, while #216's decode hot-spot fix is
retained.

Per-role thresholds (last commit, separable)

--prefill-policy and --decode-policy already choose the algorithm per role, but
every tuning knob came from a single global — parse_policy (main.rs) and
convert_policy (lib.rs) both read self.balance_*_threshold regardless of which
role they were building. There was no way to express "sticky prefill, balanced
decode". The two roles are not comparable in the first place: prefill load includes
queued requests and sits far deeper than decode, so one pair of thresholds cannot
suit both.

This adds --prefill-balance-abs-threshold, --prefill-balance-rel-threshold and
the decode equivalents, on both the CLI and the Python binding. All default to the
existing global values, so behaviour is unchanged unless the new flags are passed.

It is independent of the fix above and inert without the flags — happy to drop this
commit if you would rather keep the PR to the bug fix alone.

Checks

cargo fmt --check, cargo clippy --all-targets --all-features -- -D warnings,
cargo check --all-targets, black, ruff, and the policy and worker unit tests
all pass.

One note for CI: neither cargo test --lib nor the setuptools-rust wheel build
compiles src/main.rs, so the CLI surface is only covered by
cargo check --all-targets / cargo build.

…llm-project#197) The registry health checker unconditionally reset every worker's load counter every 10 health-check cycles, even with requests in flight. Since the cache-aware prefix tree was not reset, the policy kept routing by prefix affinity to saturated workers, creating hot spots. - Remove the periodic load reset from WorkerRegistry::start_health_checker (health checks now run in parallel) - Reset load only when a worker transitions unhealthy -> healthy, where remaining load is drift from the down period (no new requests are routed to unhealthy workers) - Delete the unused standalone start_health_checker with its conditional low-load reset (zero callers) - Redesign WorkerLoadGuard as an owning RAII guard with share()/release() and adopt it in the OpenAI router proxy path, fixing a double-decrement on retryable responses and guaranteeing paired accounting on every path - Clamp decrement at zero with a warning and wire the vllm_router_worker_load gauge on increment/decrement/reset Adds unit tests for guard semantics and registry tests proving the health checker preserves in-flight load and resets only on recovery. Signed-off-by: WU Hang <whlbx@hotmail.com>

…ssion test Address review findings on the vllm-project#197 fix: P1: resetting the load counter on an unhealthy->healthy transition can erase real in-flight load. A failed /health probe does not cancel requests routed before the failure, and between set_healthy(true) and the reset a newly routed request's increment can be wiped, or an older guard's decrement can consume a newer request's count. The health checker now never mutates load accounting; WorkerLoadGuard keeps increments and decrements paired on every request path. P2: the regression test only observed ~3 health-check cycles and could not catch the previous 10-cycle periodic reset. It now runs under paused Tokio time with a no-I/O health check and advances through 12 cycles; verified to fail against the pre-fix implementation (load reset to 0) and pass with the fix. Adds a recovery test asserting load is preserved across an unhealthy->healthy transition. Adds tokio's test-util feature as a dev-dependency for paused-time testing. Signed-off-by: WU Hang <whlbx@hotmail.com>

The only usage is the fully-qualified futures::future::join_all, which resolves via the edition extern prelude, so the use statement is unnecessary. Signed-off-by: WU Hang <whlbx@hotmail.com>

Address follow-up review findings on the vllm-project#197 fix: P1: the P/D router tracked prefill and decode load with manual increments and decrements. Cancelling the request task (e.g. client disconnect dropping the handler future) between the increment and an explicit decrement branch leaked the load permanently, since health checks no longer mask it. Both phases now use WorkerLoadGuard: the prefill guard is released at the transition to decode, the decode guard at phase completion, and cancellation releases either guard via Drop. Adds cancellation tests that run the real two-stage request flow against mock prefill/decode servers, abort the request while the upstream response is pending, and assert load returns to zero for both phases. Verified these tests fail against the previous manual accounting (load left at 1) and pass with the fix. P3: reset_load is no longer called by the runtime; update the trait and BasicWorker docs to describe it as an administrative operation that is only safe when the caller has established that no requests are in flight. Signed-off-by: WU Hang <whlbx@hotmail.com>

select_worker_with_headers gated on `max_load - min_load` across the whole fleet: if any pair of workers was imbalanced it returned early into select_worker_min_load, discarding prefix affinity for *every* request -- including requests whose preferred worker was completely idle. Under prefill/decode disaggregation that gate is essentially always open. Prefill worker load counters include queued requests, so at high concurrency the fleet spread clears balance_abs_threshold permanently, routing degenerates to shortest-queue and the prefix cache stops being consulted at all. Measured on SWE-bench Pro (4 prefill / 6 decode, 1024 rollouts in parallel, identical model and vLLM config, only the router differing): prefill prefix cache hit rate 88.9% -> 68.5%, i.e. the fraction of prompt tokens actually recomputed went 11.1% -> 31.5%; prompt tokens per rollout 675k -> 2.14M; prefill queue depth 193 -> 653; end-to-end 22.2 -> 11.5 rollouts/min. The run went from completing in 3h05 to not completing at all. This became visible only after vllm-project#197 was fixed. While the health checker was resetting load counters every 10 cycles the spread stayed near zero, so the gate never fired and routing was pure affinity by accident. Shedding load only helps when the worker holding the prefix is itself the hot one, so that is now the only case where affinity is given up: the same thresholds are compared against the selected worker rather than against the fleet maximum. select_worker_min_load had no remaining callers and is removed. Signed-off-by: Sugam Devare <sdevare@nvidia.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2b284e6b3d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-09-04T00:45:52Z

                decode_urls: final_decode_urls,
-                prefill_policy: self.prefill_policy.as_ref().map(|p| self.parse_policy(p)),
-                decode_policy: self.decode_policy.as_ref().map(|p| self.parse_policy(p)),
+                prefill_policy: self.prefill_policy.as_ref().map(|p| {


Apply role thresholds when policies fall back

In PD mode using the default --policy cache_aware without an explicit --prefill-policy, this map leaves prefill_policy as None, so RouterFactory::create_vllm_pd_router falls back to the main policy containing only the global thresholds and silently ignores --prefill-balance-*-threshold; the decode path and the mirrored Python/PyO3 conversion in src/lib.rs behave the same way. Materialize an effective role policy whenever a role-specific threshold is supplied, even if the role's policy otherwise falls back to the main policy.

Useful? React with 👍 / 👎.

…ill affinity The wheel build pinned the head of vllm-project/router#216, which fixes the decode-node hot spot (issue #197). Shipping #216 alone turns out to be worse than not shipping it for prefill-heavy benchmarks such as SWE-bench Pro. #216 makes the worker load counters honest, and that switches on a second latent bug: cache_aware decides whether to use prefix affinity from the fleet-wide load spread, so one hot worker discards affinity for every request -- including requests whose own worker is idle. Under P/D disaggregation that gate is open almost permanently, because prefill worker load counts queued requests as well as running ones. Routing degenerates to shortest-queue, already-cached prompts get recomputed, prefill saturates and decode starves behind it. Repin to the head of vllm-project/router#238, which carries #216 unchanged plus a fix applying the same load check per request, against the worker the request actually wants, rather than against the fleet maximum. Still fetched from vllm-project/router by bare SHA, as the #216 pin already was. Also in the wheel build: - VLLM_ROUTER_SRC builds from a local router checkout instead of the pinned SHA, for iterating on a fork. The directory is mounted read-only and copied into the build root so the build never writes into the source tree. - cargo check --all-targets. Neither the wheel build nor cargo test --lib compiles the vllm-router binary target, so a break in src/main.rs would ship silently. - Run the cache_aware policy tests alongside the load-accounting ones. Signed-off-by: Sugam Devare <sdevare@nvidia.com>

…ill affinity The wheel build pins the head of vllm-project/router#216, which fixes the decode-node hot spot (issue #197). Shipping #216 alone turns out to be worse than not shipping it for prefill-heavy benchmarks such as SWE-bench Pro. #216 makes the worker load counters honest, and that switches on a second latent bug: cache_aware decides whether to use prefix affinity from the fleet-wide load spread, so one hot worker discards affinity for every request -- including requests whose own worker is idle. Under P/D disaggregation that gate is open almost permanently, because prefill worker load counts queued requests as well as running ones. Routing degenerates to shortest-queue, already-cached prompts get recomputed, prefill saturates and decode starves behind it. Repin to the head of vllm-project/router#238, which carries #216 unchanged plus a fix applying the same load check per request, against the worker the request actually wants, rather than against the fleet maximum. Still fetched from vllm-project/router by bare SHA, exactly as the #216 pin already was. Signed-off-by: Sugam Devare <sdevare@nvidia.com>

…etection Moves the pinned vllm-router commit to the current head of vllm-project/router#238, which adds two failure detections on top of the prefill-affinity fix already described here. A decode worker can stop serving while still passing its health check: /health is answered by the vLLM API server, which stays responsive when the engine loop wedges. The router keeps routing to it, and the requests already dispatched are never returned. Seen on a 2 prefill / 2 decode eval as one engine holding 129 in-flight requests and going silent for 61 minutes with no health failure logged, while the run's completion rate fell to 0.1 rollouts/min. A related failure is a worker that degrades without wedging. As it stops draining, the client blocks on its outstanding requests and issues fewer new ones, so its in-flight count falls below a healthy peer's and shortest-queue routing starts to prefer it -- measured at 141 requests on an engine serving 1.0 tok/s per sequence against 105 on a peer at 49.6. Both are detected from per-worker completion progress rather than liveness alone, and both feed the router's existing failure-threshold machinery, so a flagged worker stops receiving traffic and recovers on its own once it makes progress again. Signed-off-by: Sugam Devare <sdevare@nvidia.com>

… fixes (#3080) ## What does this PR do? Repins the `vllm-router` build used by the Super 3.5 eval container, from the head of [vllm-project/router#216](vllm-project/router#216) to the head of [#238](vllm-project/router#238). #216 fixes the decode-node hot spot (issue #197), but shipping it alone is worse than not shipping it for prefill-heavy benchmarks like SWE-bench Pro. #216 makes the worker load counters honest, and that switches on a second latent bug: `cache_aware` decides whether to use prefix affinity from the **fleet-wide** load spread, so one hot worker discards affinity for *every* request — including requests whose own worker is idle. Under P/D disaggregation that gate is open almost permanently, because prefill worker load counts queued requests as well as running ones. Routing degenerates to shortest-queue, already-cached prompts get recomputed, prefill saturates, and decode starves behind it. #238 carries #216 unchanged plus a fix that applies the same load check per request, against the worker the request actually wants, rather than against the fleet maximum. Still fetched from `vllm-project/router` by bare SHA, exactly as the #216 pin already was — a PR head is a ref in the upstream repo even when the branch lives on a contributor's fork. ## Testing Validated on SWE-bench Pro (4 prefill / 6 decode) against two baselines on the same checkpoint and vLLM config: the pre-#216 router, and #216 alone. With this pin the prefill prefix-cache hit rate and end-to-end throughput return to pre-#216 behaviour while keeping #216's decode hot-spot fix. Signed-off-by: Sugam Devare <sdevare@nvidia.com>

#3107) The pin added in #3080 pointed at a commit that also carried per-role balance threshold flags (--prefill-balance-abs-threshold and friends). They have been dropped from vllm-project/router#238 to keep the PR to the bug fix, so this moves the pin to the new head. No behaviour change. The remaining commits are unchanged: #216 (health checker no longer resets in-flight worker load) plus the fix for the prefill-side regression #216 exposes, where cache_aware decided prefix affinity from the fleet-wide load spread instead of per request. Signed-off-by: Sugam Devare <sdevare@nvidia.com>

… fixes (NVIDIA-NeMo#3080) ## What does this PR do? Repins the `vllm-router` build used by the Super 3.5 eval container, from the head of [vllm-project/router#216](vllm-project/router#216) to the head of [NVIDIA-NeMo#238](vllm-project/router#238). NVIDIA-NeMo#216 fixes the decode-node hot spot (issue NVIDIA-NeMo#197), but shipping it alone is worse than not shipping it for prefill-heavy benchmarks like SWE-bench Pro. NVIDIA-NeMo#216 makes the worker load counters honest, and that switches on a second latent bug: `cache_aware` decides whether to use prefix affinity from the **fleet-wide** load spread, so one hot worker discards affinity for *every* request — including requests whose own worker is idle. Under P/D disaggregation that gate is open almost permanently, because prefill worker load counts queued requests as well as running ones. Routing degenerates to shortest-queue, already-cached prompts get recomputed, prefill saturates, and decode starves behind it. NVIDIA-NeMo#238 carries NVIDIA-NeMo#216 unchanged plus a fix that applies the same load check per request, against the worker the request actually wants, rather than against the fleet maximum. Still fetched from `vllm-project/router` by bare SHA, exactly as the NVIDIA-NeMo#216 pin already was — a PR head is a ref in the upstream repo even when the branch lives on a contributor's fork. ## Testing Validated on SWE-bench Pro (4 prefill / 6 decode) against two baselines on the same checkpoint and vLLM config: the pre-NVIDIA-NeMo#216 router, and NVIDIA-NeMo#216 alone. With this pin the prefill prefix-cache hit rate and end-to-end throughput return to pre-NVIDIA-NeMo#216 behaviour while keeping NVIDIA-NeMo#216's decode hot-spot fix. Signed-off-by: Sugam Devare <sdevare@nvidia.com> Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com>

…DIA-NeMo#238 (NVIDIA-NeMo#3107) The pin added in NVIDIA-NeMo#3080 pointed at a commit that also carried per-role balance threshold flags (--prefill-balance-abs-threshold and friends). They have been dropped from vllm-project/router#238 to keep the PR to the bug fix, so this moves the pin to the new head. No behaviour change. The remaining commits are unchanged: NVIDIA-NeMo#216 (health checker no longer resets in-flight worker load) plus the fix for the prefill-side regression NVIDIA-NeMo#216 exposes, where cache_aware decided prefix affinity from the fleet-wide load spread instead of per request. Signed-off-by: Sugam Devare <sdevare@nvidia.com> Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com>

… fixes (#3080) ## What does this PR do? Repins the `vllm-router` build used by the Super 3.5 eval container, from the head of [vllm-project/router#216](vllm-project/router#216) to the head of [#238](vllm-project/router#238). #216 fixes the decode-node hot spot (issue #197), but shipping it alone is worse than not shipping it for prefill-heavy benchmarks like SWE-bench Pro. #216 makes the worker load counters honest, and that switches on a second latent bug: `cache_aware` decides whether to use prefix affinity from the **fleet-wide** load spread, so one hot worker discards affinity for *every* request — including requests whose own worker is idle. Under P/D disaggregation that gate is open almost permanently, because prefill worker load counts queued requests as well as running ones. Routing degenerates to shortest-queue, already-cached prompts get recomputed, prefill saturates, and decode starves behind it. #238 carries #216 unchanged plus a fix that applies the same load check per request, against the worker the request actually wants, rather than against the fleet maximum. Still fetched from `vllm-project/router` by bare SHA, exactly as the #216 pin already was — a PR head is a ref in the upstream repo even when the branch lives on a contributor's fork. ## Testing Validated on SWE-bench Pro (4 prefill / 6 decode) against two baselines on the same checkpoint and vLLM config: the pre-#216 router, and #216 alone. With this pin the prefill prefix-cache hit rate and end-to-end throughput return to pre-#216 behaviour while keeping #216's decode hot-spot fix. Signed-off-by: Sugam Devare <sdevare@nvidia.com> Signed-off-by: Junkeun Yi <jkyi@nvidia.com>

#3107) The pin added in #3080 pointed at a commit that also carried per-role balance threshold flags (--prefill-balance-abs-threshold and friends). They have been dropped from vllm-project/router#238 to keep the PR to the bug fix, so this moves the pin to the new head. No behaviour change. The remaining commits are unchanged: #216 (health checker no longer resets in-flight worker load) plus the fix for the prefill-side regression #216 exposes, where cache_aware decided prefix affinity from the fleet-wide load spread instead of per request. Signed-off-by: Sugam Devare <sdevare@nvidia.com> Signed-off-by: Junkeun Yi <jkyi@nvidia.com>

wuhang2014 and others added 5 commits August 16, 2026 15:53

chore: remove redundant single-component futures import

ad83cec

The only usage is the fully-qualified futures::future::join_all, which resolves via the edition extern prelude, so the use statement is unnecessary. Signed-off-by: WU Hang <whlbx@hotmail.com>

chatgpt-codex-connector Bot reviewed Sep 4, 2026

View reviewed changes

sdevare-nv force-pushed the sdd/cache-aware-per-request-affinity branch from 2b284e6 to 5b7ce08 Compare September 4, 2026 00:45

sdevare-nv changed the title ~~Sdd/cache aware per request affinity~~ cache aware per request affinity Sep 4, 2026

sdevare-nv changed the title ~~cache aware per request affinity~~ feat: cache aware per request affinity Sep 4, 2026

sdevare-nv mentioned this pull request Sep 4, 2026

fix(nemotron_3.5_super): pin a vllm-router build with prefix affinity fixes NVIDIA-NeMo/Gym#3080

Merged

sdevare-nv force-pushed the sdd/cache-aware-per-request-affinity branch from 11f6463 to 5b7ce08 Compare September 4, 2026 16:12

sdevare-nv force-pushed the sdd/cache-aware-per-request-affinity branch from 5b7ce08 to 84b4425 Compare September 4, 2026 21:37

sdevare-nv mentioned this pull request Sep 4, 2026

fix(nemotron_3.5_super): repin vllm-router to the current head of #238 NVIDIA-NeMo/Gym#3107

Merged

Navigation Menu

Uh oh!

feat: cache aware per request affinity - #238

feat: cache aware per request affinity#238
sdevare-nv wants to merge 5 commits into
vllm-project:mainfrom
sdevare-nv:sdd/cache-aware-per-request-affinity

sdevare-nv commented Sep 4, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

sdevare-nv commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

The problem

The fix

Per-role thresholds (last commit, separable)

Checks

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sdevare-nv commented Sep 4, 2026 •

edited

Loading