Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a cache-aware routing mechanism called PREFIX_AFFINITY to the DataParallelController. This method routes requests sharing a prefix to the same data-parallel rank using rendezvous (HRW) hashing over live ranks, while using a load guard to prevent hotspots. It also adds configuration arguments and comprehensive unit tests. The review feedback suggests several performance optimizations: leveraging .tobytes() for faster token prefix hashing, calculating the overload threshold once outside the loop in the scheduler, caching encoded rank bytes in _rendezvous_ranked to avoid redundant string formatting, and removing the unused _rank_is_overloaded helper method.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| for token in input_ids[:prefix_len]: | ||
| h.update(int(token).to_bytes(8, "little", signed=True)) | ||
| return "token-prefix:" + h.hexdigest() |
There was a problem hiding this comment.
Performance Optimization: Slicing input_ids and calling to_bytes in a loop for every token up to prefix_len (default 4096) is extremely slow in Python. Since input_ids is typically an array object, we can leverage its high-performance .tobytes() method to convert the entire slice to bytes in a single fast C-level memory copy, falling back to the loop only for non-array types (like lists in unit tests). This can speed up token prefix hashing by up to 400x.
| for token in input_ids[:prefix_len]: | |
| h.update(int(token).to_bytes(8, "little", signed=True)) | |
| return "token-prefix:" + h.hexdigest() | |
| h = hashlib.blake2b(digest_size=16) | |
| sliced = input_ids[:prefix_len] | |
| if hasattr(sliced, "tobytes"): | |
| h.update(sliced.tobytes()) | |
| else: | |
| for token in sliced: | |
| h.update(int(token).to_bytes(8, "little", signed=True)) |
There was a problem hiding this comment.
Done in 36d2ee9 — added a .tobytes() fast path for array-like input_ids, falling back to the per-token loop for plain lists.
| live = self._live_ranks() | ||
| for rank in self._rendezvous_ranked(route_key, live): | ||
| if not self._rank_is_overloaded(rank): | ||
| self._increment_rank_budget(rank, req) | ||
| self.workers[rank].send_pyobj(req) | ||
| return |
There was a problem hiding this comment.
Performance Optimization: The average load of the live ranks is constant for the duration of a single scheduler call. Currently, _rank_is_overloaded is called inside the loop over all ranked candidates, which redundantly recalculates the average load of all live ranks up to
| live = self._live_ranks() | |
| for rank in self._rendezvous_ranked(route_key, live): | |
| if not self._rank_is_overloaded(rank): | |
| self._increment_rank_budget(rank, req) | |
| self.workers[rank].send_pyobj(req) | |
| return | |
| live = self._live_ranks() | |
| if len(live) <= 1: | |
| overload_threshold = float("inf") | |
| else: | |
| loads = [self._rank_load(r) for r in live] | |
| avg = max(sum(loads) / len(loads), 1.0) | |
| overload_threshold = self._affinity_max_load_skew * avg | |
| for rank in self._rendezvous_ranked(route_key, live): | |
| if self._rank_load(rank) <= overload_threshold: | |
| self._increment_rank_budget(rank, req) | |
| self.workers[rank].send_pyobj(req) | |
| return |
There was a problem hiding this comment.
Done in 36d2ee9 — the overload threshold is now computed once per dispatch before the candidate loop instead of per candidate.
There was a problem hiding this comment.
Note: I kept the dispatch call as sock_send(self.workers[rank], req) rather than the raw send_pyobj from the suggested snippet, to stay consistent with the current controller convention (per an earlier review comment).
| def _rank_is_overloaded(self, rank: int) -> bool: | ||
| """True if ``rank``'s load exceeds ``max_load_skew`` times the average.""" | ||
| live = self._live_ranks() | ||
| if len(live) <= 1: | ||
| return False | ||
| loads = [self._rank_load(r) for r in live] | ||
| avg = max(sum(loads) / len(loads), 1.0) | ||
| return self._rank_load(rank) > self._affinity_max_load_skew * avg |
There was a problem hiding this comment.
Clean Code / Maintainability: Since the overload check has been optimized to run inline and calculate the average load once per scheduler call, the private helper method _rank_is_overloaded is no longer used and can be safely removed to keep the codebase clean.
There was a problem hiding this comment.
Done in 36d2ee9 — inlined the threshold check into the scheduler loop and removed the now-unused _rank_is_overloaded helper.
| def _rendezvous_ranked(self, key: str, ranks: List[int]) -> List[int]: | ||
| """Rendezvous (HRW) order of ``ranks`` for ``key``, best score first.""" | ||
| key_b = key.encode("utf-8", errors="surrogatepass") | ||
| return sorted( | ||
| ranks, | ||
| key=lambda rank: self._hash64(key_b + b"\x00" + str(rank).encode("ascii")), | ||
| reverse=True, | ||
| ) |
There was a problem hiding this comment.
Performance Optimization: In _rendezvous_ranked, formatting and encoding the rank integers (str(rank).encode("ascii")) inside the sorting key lambda is executed repeatedly for every rank comparison. Pre-populating a lazy cache of the encoded rank bytes once before the sorted call avoids this overhead entirely.
| def _rendezvous_ranked(self, key: str, ranks: List[int]) -> List[int]: | |
| """Rendezvous (HRW) order of ``ranks`` for ``key``, best score first.""" | |
| key_b = key.encode("utf-8", errors="surrogatepass") | |
| return sorted( | |
| ranks, | |
| key=lambda rank: self._hash64(key_b + b"\x00" + str(rank).encode("ascii")), | |
| reverse=True, | |
| ) | |
| def _rendezvous_ranked(self, key: str, ranks: List[int]) -> List[int]: | |
| """Rendezvous (HRW) order of ``ranks`` for ``key``, best score first.""" | |
| key_b = key.encode("utf-8", errors="surrogatepass") | |
| encoded = self.__dict__.setdefault("_encoded_ranks", {}) | |
| for r in ranks: | |
| if r not in encoded: | |
| encoded[r] = b"\x00" + str(r).encode("ascii") | |
| return sorted( | |
| ranks, | |
| key=lambda rank: self._hash64(key_b + encoded[rank]), | |
| reverse=True, | |
| ) |
There was a problem hiding this comment.
Done in 36d2ee9 — each rank's HRW score is precomputed into a dict once, then used as the sort key, so the encode/hash runs once per rank instead of on every comparison.
edaa7fb to
36d2ee9
Compare
|
Could someone with CI permission please add the |
|
Hi @hnyls2002 The main runtime change is in The feature is fully opt-in, reuses the existing Thanks! |
|
Been chasing the same thing on AMD and ended up at basically this design, so I'd rather help land #31170 than push a competing PR. On 8x MI355X with GLM5.2-FP8, dp8/tp8, Some numbers from our runs in case they help:
One thing, more a question than a blocker: does Happy to help review, and I can pull more detail off the MI355X runs if useful. |
|
@reger-men Nice to see that we arrived at the same design! |
|
Hi @hnyls2002 @merrymercy @Ying1123 @xiezhq-hermann Thanks! |
bfd8373 to
fa56573
Compare
|
Thanks for putting this together. I encountered the same affinity-vs-hotspot trade-off at the SGLang Model Gateway's inter-worker routing layer and opened #33625 to discuss an opt-in The scopes are complementary: this PR handles affinity and overload spillover among native DP ranks within a single SGLang instance, while #33625 concerns routing among separate backend workers at the gateway. I have aligned the gateway proposal with this PR's I am not proposing a competing native-DPA implementation. If helpful, I would be glad to contribute the end-to-end entrypoint test mentioned in this PR, verifying that |
…project#32035 sgl-project#33656 sgl-project#32183 sgl-project#33145) Applied PRs (latest from GitHub): sgl-project#33288 Indexer logits OOM fix sgl-project#30393 HiCache packed/sidecar draft caches sgl-project#31170 DPA prefix_affinity load balancing sgl-project#33795 DSpark compact ragged-verify CUDA graph JIT race sgl-project#32467 C128 plan-kernel warp barrier sgl-project#33865 DSpark x prefill CP unblock sgl-project#30371 SWA state pool sizing (storage page) sgl-project#33358 FlashMLA norm-rope K-tokens-per-block ILP sgl-project#33872 num_draft_tokens clamp + extend_len==0 skip (supersede sgl-project#32183) sgl-project#34002 Sidecar backup vacuously-successful fix (replaces sgl-project#33656, with tests) sgl-project#33862 Reclaim redundant host mirrors after storage backup sgl-project#31315 Avoid repeated Mooncake gets after stale hits sgl-project#32327 Q8KV8 sparse MLA prefill backend (flashmla_sparse_q8) sgl-project#31668 Fix sidecar pool life-time (use-after-free on prefetch abort) sgl-project#31195 TP0 verify-token-budget broadcast (adapted to get_schedule() API) Dropped (per user request or superseded): sgl-project#32771 IndexCache C4 top-k reuse — has bug sgl-project#32035 DSpark C128 online compressor — has bug sgl-project#33656 Superseded by sgl-project#34002 (same fix + unit tests) sgl-project#32183 Superseded by sgl-project#33872 (included in supersede PR) sgl-project#33145 Base f01f706 already has superior reasoning-effort profile system Conflicts resolved: sgl-project#31195: adapted to base get_schedule().disable_overlap_schedule API sgl-project#32327: path remapped jit_kernel/ -> kernels/jit/ and kernels/ops/attention/ sgl-project#31668: applied cleanly on top of sgl-project#30393+sgl-project#34002+sgl-project#33862 modifications
|
Backported this to v0.5.16 for a DP=4/EP=4 DeepSeek-V4-Flash deployment on SM120 and ran it in production for a few days. Wanted to flag a real issue in the overload guard's load signal, since it reproduces on this PR as written (not just in our backport).
def _rank_load(self, rank: int) -> float:
if self._affinity_fallback == LoadBalanceMethod.TOTAL_REQUESTS:
return float(self.dp_budget.total_requests[rank])
return float(self.dp_budget.total_tokens[rank])Under real traffic with long-lived, large-context sessions, this misfires: a rank that's steadily decoding one 300K–800K-token conversation reads as many times the fleet-average token footprint, even though it has plenty of compute/memory headroom to keep serving. The skew guard then treats it as overloaded and reroutes that same session's next turn to a different rank — which evicts exactly the radix-cache prefix affinity exists to preserve, forcing a full re-prefill of the whole context. We measured this directly in prod: the large majority of prefill starts were landing with The fix that worked for us: key the guard off live running+waiting request count instead, decoupled from def _rank_load(self, rank: int) -> float:
return float(self.dp_budget.total_requests[rank])
This also matches the pattern in every other cache-affinity + overload-protection system we looked at while chasing this down: Happy to open this as a follow-up PR against your branch, or fold it in here directly if that's easier — whichever you'd prefer. We also extended the unit tests with a case pinning this exact scenario (huge resident-token footprint, zero requests, on the preferred rank → affinity should hold) if useful. |
…/session affinity
fa56573 to
4ce74f3
Compare
|
@hassellof Great findings, ty! |
7d136aa to
6fc452a
Compare
|
@ispobock updated branch to latest main, conflicts are resolved. Could you please run /tag-and-rerun-ci when convenient? |
…alancing
Single-instance DP attention only offers content-agnostic load balancing
(round_robin/total_requests/total_tokens), so a shared prefix only hits a
rank's radix cache about once per dp_size requests. This adds a new
--load-balance-method prefix_affinity that routes by a session-level key:
1. req.routing_key (e.g. x-smg-routing-key) when present, so callers can
pin affinity at the agent/session/project level;
2. a hash of the first N input tokens as a fallback (disable-able).
Routing uses rendezvous (HRW) hashing over the live ranks so a rank drop or
dp_size change only remaps the affected keys, and an overload guard walks the
HRW-ranked candidates to skip ranks above max_load_skew x the fleet average,
degrading gracefully instead of hotspotting one rank under concurrency.
Opt-in; reuses the existing routing_key plumbing (no protocol/io_struct
changes). Config: --prefix-affinity-fallback / --prefix-affinity-max-load-skew
/ --prefix-affinity-hash-tokens / --prefix-affinity-disable-token-fallback.
Adds CPU unit tests covering affinity, distinct-key spread, HRW stability,
the overload-guard spill, and the fallback paths.
TokenizedEmbeddingReqInput does not declare routing_key (unlike TokenizedGenerateReqInput) but still reaches prefix_affinity_scheduler via dispatching_with_trace, so a bare req.routing_key raised AttributeError on the embedding path. Read it via getattr so such requests fall through to the token/load fallback. Adds a CPU unit test driving a routing_key-less req through the scheduler.
- Guard _token_prefix_key and _estimated_tokens against multi-element NumPy array truthiness errors (if not input_ids raises ValueError). Use explicit None/len() checks instead. - Fix prefix_affinity_fallback CLI help: a missing routing key goes through token-prefix affinity first; the load-aware fallback only runs when token fallback is disabled or no token key is available. - Black-format the new test assertion.
_rank_load used total_tokens (resident KV) when _affinity_fallback != TOTAL_REQUESTS (the default), so a rank decoding one long-context session (300K+ tokens) read as overloaded and got its next turn rerouted — evicting the exact radix-cache prefix affinity preserves. Reported by @hassellof from a multi-day DeepSeek-V4-Flash production deployment (DP=4/EP=4, SM120) where the majority of prefill starts landed with #cached-token: 0. Fix: always use total_requests (num_running_reqs + num_waiting_reqs), decoupled from _affinity_fallback. Matches sgl-router and other cache-affinity routers that key imbalance off active request count. Updated the existing spill test to use total_requests, added a regression test for the huge-token-footprint-zero-requests case.
c9c6784 to
381212e
Compare
|
Would be cool to get this merged. I can try it on GLM-5.3 with a B200 node if that helps |
7 participants
TL;DR
prefix_affinityis an opt-in, routing-key-based affinity policy for single-instance DPattention. It targets multi-turn / agentic workloads where independent sessions share a large
system/tools prefix but should not all be routed to the same DP rank. Compared with
round_robin, it preserves per-session radix-cache locality; compared with first-N-token prefixhashing (#26612), it avoids collapsing all sessions whose distinguishing tokens appear after the
shared prefix window. An overload guard lets affinity degrade to load-aware routing when the
preferred rank becomes hot.
Motivation
Single-instance DP attention (
--enable-dp-attention --dp-size N) currently offers onlyround_robin,follow_bootstrap_room,total_requests, andtotal_tokens— none route byrequest content, so under
round_robina repeated shared prefix only hits a rank's radixcache about once per
dp_sizerequests. This is the known DPA cache-miss problem.The open proposal #26612 (
prefix_match) addresses the sequential case by hashing the first~4K tokens and routing
hash % dp_size. This PR covers an adjacent failure mode: for workloadswhere the first hash window is dominated by a shared system/tools prompt (agentic / OpenCode /
Claude-Code-style traffic), first-N-token hashing cannot distinguish sessions. In the worst case,
independent sessions with identical leading tokens map to the same DP rank, creating a hotspot.
This policy was motivated by our production-like GLM5.2 agentic workload, where many concurrent local-agent sessions share a large system/tools prefix. Existing native-DPA policies either miss cache locality (
round_robin) or scatter session turns under concurrency (total_tokens), while first-N-token prefix hashing cannot distinguish sessions when the shared prefix dominates the hash window.Why not just use SMG?
SMG is the right layer for production inter-instance routing. This PR fixes a different layer:
intra-instance routing among native DPA ranks. When a single SGLang server is launched with
--enable-dp-attention --dp-size N, requests are dispatched internally byDataParallelController; an external gateway cannot choose the internal DP rank. Therefore SMGdoes not address the native-DPA radix-cache locality problem unless users give up this topology
and run separate replicas instead.
prefix_affinityis a small opt-in policy for users who neednative DPA's KV-capacity benefits and also need per-session cache locality.
What this PR adds
A new
--load-balance-method prefix_affinityfor the single-instanceDataParallelControllerthat routes by a session-level key and preserves affinity under load instead of collapsing:
Routing key priority — reuses the existing
req.routing_keyas the primary affinitykey so callers can pin at the agent/session/project level. This PR does not introduce a new
public request field: it consumes
req.routing_key, which the native OpenAI-compatibleentrypoints already populate from the
x-smg-routing-keyheader viaOpenAIServingBase.extract_routing_key(serving_completions.py,serving_chat.py, …) —so it works on native single-instance DPA without SMG in front.
Selection order:
req.routing_keyis present, use HRW affinity on that key.first
--prefix-affinity-hash-tokenstokens.--prefix-affinity-fallback.Rendezvous (HRW) hashing instead of
hash % dp_size: candidates are scored using thestable DP rank ids, not their position in a filtered live-rank list, so taking an offline
rank out of the candidate set does not renumber the others. When a rank drops or
dp_sizechanges, only keys that mapped to the affected rank move — everyone else keeps affinity.
Overload guard — walks the HRW-ranked candidates and picks the best rank whose active
request count is within
max_load_skew ×the average across live ranks. The guard usesdp_budget.total_requests(running + waiting) rather than resident token/KV footprint,so a long-context session is not treated as overloaded merely because its useful cache is
large. The configured fallback policy remains independent.
New flags
--load-balance-method prefix_affinity--prefix-affinity-fallback {round_robin,total_requests,total_tokens}total_tokens--prefix-affinity-max-load-skew FLOAT (>=1.0)1.5--prefix-affinity-hash-tokens INT4096--prefix-affinity-disable-token-fallbackBlast radius / compatibility
managers/data_parallel_controller.py(new enum value +scheduler method + helpers) and
server_args.py(choice + 4 config fields + validation).test/registered/unit/managers/test_data_parallel_controller.py.req.routing_key.budget accounting/status handling are identical.
Benchmarks
8×H100 80GB, sglang 0.5.15 + this feature.
gsp= shared-prefix agentic generator with--gsp-send-routing-key. Baselineround_robin== current upstream DPA behavior.1.
prefix_affinityvsround_robinvstotal_tokens— the PR's core A/B (Qwen3.5-122B-A10B-FP8)16 groups × 32 prompts, 4096-tok system prompt, q128/o256, 512 prompts, concurrency 64. Same seed, same server args except
--load-balance-method.Win on the targeted DPA workload — this is the contribution. The relevant native-DPA
baseline is
round_robinfor the common single-instance non-PD setup.total_tokensis the other obviousbaseline a reviewer may ask for: per issue #26611,
total_tokenscan accidentally approximatelocality under strictly sequential traffic but falls apart under concurrency — here it is the
worst of the three (1.36 req/s, ~4× slower than
round_robin), because load-following with nocontent awareness scatters each session's turns across ranks and thrashes every rank's radix
cache.
prefix_affinitybeats both by keeping each routing-key's turns resident on one rank whilethe overload guard prevents hotspotting.
2. KV cache freed by DP attention (why DPA + this routing matters)
3. DPA +
prefix_affinityvs TP-pure — DeepSeek-V4-Flash-FP8 (MLA) showcase64 groups × 8 prompts, 8192-tok prefix, q256/o512, 512 prompts, concurrency 128.
MLA has a single KV head, so pure TP replicates the full KV cache on every rank; DP
attention gives each rank its own pool (6.5× aggregate KV), and routing-key affinity keeps
shared-prefix sessions resident → improves all reported metrics for this workload.
Honest scope note
For MoE models with small active params (Qwen), DPA frees KV and improves TTFT under
pressure but loses aggregate throughput/TPOT vs TP-pure, because each request decodes on
1/
dp_sizeof the hardware. The DPA-vs-TP latency win is strongest for MLA models(DeepSeek). Within DPA, the claim this PR makes is narrow and defensible: for
repeated-prefix / multi-turn / agentic workloads where callers provide stable routing keys,
prefix_affinitybeatsround_robinin the DPA benchmarks reported here. It is not claimed tohelp random, non-repeating workloads. Without an explicit routing key, it uses token-prefix
affinity by default, or the configured load-aware fallback when token fallback is disabled or no
token key is available.
Usage
For workloads with a large shared system/tools prompt, callers should provide a stable
per-session or per-project routing key. Without a routing key, token fallback behaves like prefix
hashing and may not distinguish sessions whose first N tokens are identical.
Tests
Added to
test/registered/unit/managers/test_data_parallel_controller.py(CPU CI, no GPU),following the existing
__new__-fixture convention. 27/27 pass (14 pre-existing + 13 new):routing_keyselects the same HRW-preferred rank under balanced load(affinity); 256 distinct keys sharing an identical 8K prefix distribute across all 8 live ranks
in the test setup (the anti-collapse fix vs dp-attention: add prefix_match load balance for in-instance cache-aware routing #26612's first-N-token
hash); a hot shared key under accumulating load spills across ranks via the overload guard
instead of hotspotting one, while the first low-load request still lands on the HRW-preferred rank.
unlike
hash % dp_size.candidate (and records the dispatch in the budget); a single live rank never deadlocks.
--prefix-affinity-disable-token-fallbackdefers to the load-aware fallback method; an explicit
routed_dp_rankbypasses affinity entirely;a request without a
routing_keyattribute (e.g.TokenizedEmbeddingReqInputon the embeddingpath) falls through to the token/load fallback instead of raising.
The
x-smg-routing-key→req.routing_keypath is pre-existing plumbing(
OpenAIServingBase.extract_routing_key, wired intoserving_completions/serving_chat/…), notintroduced here, so the unit tests inject
routing_keydirectly on the request. An end-to-endentrypoint test (assert the header reaches
req.routing_keywithout SMG) can be added ifreviewers want belt-and-suspenders coverage.
Local validation
pre-commit run --all-files(black / isort / ruff / codespell) — clean on the changed filespytest test/registered/unit/managers/test_data_parallel_controller.py -qCI States
Latest PR Test (Base): ❌ Run #33373193048
Latest PR Test (Extra): ❌ Run #33373192982
Latest PR Test (AMD ROCm 7.2): ❌ Run #33373192966