feat(inference): prefix-affinity routing and a configurable cost policy in the DP coordinator - #6668
Draft
sidsingh-nvidia wants to merge 9 commits into
Draft
Conversation
The replicas shared one listening socket: the parent bound it and passed the same fd to every forked worker, so all of them accepted from a single queue. That does not balance. Whichever worker is already running tends to win the wakeup, and it keeps winning, because an event loop with work in flight polls more often than one blocked in accept. SO_REUSEPORT was set on that socket but was inert -- the kernel only load-balances when several sockets are bound to the port and it can hash a connection's 4-tuple to choose between them. Measured with 32 replicas and a fresh connection per request, ~90% of traffic landed on 5 of them, 20 replicas served exactly one request each, and throughput was 3.7x lower than the same server under a pooled client that opened its connections up front. Load made it worse rather than averaging it out: at 2048 requests the busiest replica took 604x the quietest. Each replica now binds its own socket on the shared port, so every one gets its own accept queue. Spread became max/min 1.5x with all replicas serving, and throughput 2.5x on the fresh-connection path. A pooled client is unaffected in steady state, which is the point: how well the frontend spreads no longer depends on connection behaviour the server cannot observe. start_text_gen_server now returns the base URL it is serving on. Callers that start a frontend on more than one rank need the addresses to spread requests over; previously they had to reconstruct them. The signature is otherwise unchanged, including sock, which still fixes the port -- it is closed rather than shared, since replicas bind their own. tools/run_dynamic_text_generation_server.py gains --frontend-on-all-ranks, which hosts a frontend on every rank and gathers the URLs. Frontend work is CPU-bound and otherwise confined to one rank's CPU allocation while the rest of the job's cores go unused. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com> (cherry picked from commit 1001ca9)
…is metadata of constant size. This is the only thing that the coordinator needs to unpack/read and pack
Route each request to the rank that already holds the most of its prompt rather than to the least loaded one. Frontends compute per-block prompt hashes and ship them alongside the request; the coordinator scores ranks by (prefill blocks still to compute) x (1 + load), so idle ranks fill first and a rank holding the prefix wins thereafter. Hashing happens on the frontend, not the coordinator: the tokens are already in hand there, frontends run many-to-one against a single serial coordinator loop, and hashing at the coordinator would mean unpacking the prompt frame the request/prompt frame split exists to avoid. Whether to hash follows from the coordinator's routing policy, which is the only component that knows if anyone will read the hashes. `routes_on_prefix` lives beside the policy enum so a new prefix-aware policy is a one-line change. `block_size_tokens` is granularity only and is always passed; it must match the engine's KV block size or the hashes name blocks the engine never cached. Coordinator-side cache tracking assumes an engine still holds a block for `prefix_cache_ttl_seconds` under the LRU eviction policy, since engine-side eviction is not observable from the coordinator. On a 16-engine nanoV3.5 SWE-RL run this moved the prefill skip rate from 59-63% to 97.7% and the prefill share of step time from 42-62% to 10.3%. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com> (cherry picked from commit fcb23e1)
The coordinator conflated two decisions: which affinity signal to read
(first-block hit vs contiguous prefix depth) and how that affinity is weighed
against rank load. Split them so the cost function is selectable independently
of the signal.
`PrefixCachingCoordinatorPolicy` keeps its meaning -- it picks the signal --
and both signals are now normalized to the fraction of the request already
cached on a rank, in [0, 1]. The new `PrefixCachingCostPolicy` picks how that
fraction is scored, and composes with either signal:
RELATIVE_LOAD_WEIGHTED (new default)
score = fraction - beta * (load - mean) / max(1, mean)
Load is measured against the fleet mean, so the penalty vanishes while
ranks are balanced: at saturation this is pure affinity, and load only
pulls toward idle ranks as the fleet diverges. Approximates the session
stickiness a session-affinity router gets for free, with no session id.
The mean is floored at 1 so a near-idle fleet does not turn one in-flight
request into a large relative load and thrash on noise.
FREE_CAPACITY_WEIGHTED (the previous first_prefix_block behaviour)
score = alpha * fraction + (1 - alpha) * free_slots / max_requests
Fixes the trade-off in absolute terms rather than relative to load.
Also flips the coordinator policy default from LOAD_BALANCED to LONGEST_PREFIX,
so prefix affinity is used by default once prefix caching is enabled.
Measured on a 16-replica SWE rollout workload, RELATIVE_LOAD_WEIGHTED cuts
drain-phase imbalance sharply versus the previous multiplicative cost:
busy-phase load CV 0.434 -> 0.272, tail CV 1.275 -> 0.555,
idle-while-others-busy 15.7% -> 6.7%, max/median replica load 2.13 -> 1.44.
Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
svcnvidia-nemo-ci
marked this pull request as draft
August 19, 2026 18:12
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
6 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Two things, split so the second is reviewable on its own:
The split
The coordinator previously conflated two decisions: which affinity signal to read and how that signal trades off against load. They are now orthogonal.
PrefixCachingCoordinatorPolicykeeps its meaning and picks the signal. Both prefix-aware signals are normalized to the fraction of the request already cached on a rank, in[0, 1]:LONGEST_PREFIXFIRST_PREFIX_BLOCKLOAD_BALANCEDThe new
PrefixCachingCostPolicypicks how that fraction is scored, and composes with either prefix-aware signal:RELATIVE_LOAD_WEIGHTED(default)fraction - beta * (load - mean) / max(1, mean)FREE_CAPACITY_WEIGHTEDalpha * fraction + (1 - alpha) * free_slots / max_requestsFREE_CAPACITY_WEIGHTEDis the pre-existingfirst_prefix_blockbehaviour, preserved under a name.Why RELATIVE_LOAD_WEIGHTED is the default
Measuring load against the fleet mean rather than in absolute terms is what makes the penalty behave correctly at both ends:
Both terms are normalized, so
betais dimensionless:beta = 1.0means a rank at twice the mean forfeits one full prompt's worth of cache credit. The mean is floored at 1 so a near-idle fleet does not turn a single in-flight request into a large relative load and thrash on noise.This approximates the stickiness a session-affinity router gets for free — a multi-turn request lands back on the rank holding its history — without needing a session id. The prior multiplicative cost instead charged
remaining_blocks * (1 + load), which over-weights load in the tail.Measured impact
16-replica agentic (SWE) rollout workload,
RELATIVE_LOAD_WEIGHTEDvs. the previous multiplicative cost:Prefix-cache skip rate on the same workload rose from 59-63% to 97.7%.
Behaviour matrix
Two replicas, A holds the full prefix, B is colder:
relative_load_weightedfree_capacity_weighted[32, 32][20, 10][8, 2][4, 0][0, 0]The default diverges only where one rank is meaningfully idle relative to the fleet.
Default change — please flag if this is contentious
prefix_caching_coordinator_policyflips fromLOAD_BALANCEDtoLONGEST_PREFIX, so prefix affinity is used by default once prefix caching is enabled. Runs with prefix caching off are unaffected. Happy to land this behind the old default if reviewers would rather stage it.New flags
--inference-dynamic-batching-prefix-caching-cost-policy(defaultrelative_load_weighted)--inference-dynamic-batching-prefix-caching-load-beta(default1.0)--inference-dynamic-batching-prefix-caching-routing-alphais unchanged but now documented as applying only underfree_capacity_weighted.