Skip to content

fix(pd): balance prefill nodes on decayed recent load to avoid cache-affinity starvation - #1433

Open
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:fix/cache-aware-prefill-starvation
Open

fix(pd): balance prefill nodes on decayed recent load to avoid cache-affinity starvation#1433
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:fix/cache-aware-prefill-starvation

Conversation

@sufubao

@sufubao sufubao commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

问题

LoadBalancedCacheAwareSelector 的 1.2× 均衡门闩用的是 PD_Client_Obj.dispatched_prompt_chars,即自启动起单调累计的派发字符数。在多 prefill 节点 + cache 亲和场景下存在单点饥饿:

  • 同前缀的请求会被 cache 亲和持续钉在某个 P 节点;
  • 另一个 P 节点虽然已经空闲,但其历史累计派发量仍然很高
  • 门闩判断 max/min < 1.2 永远成立,从不触发,空闲节点被饿死。

复现

m39、Qwen3.5-4B、2P1D 拓扑,40 个相同 prompt 串行打:

指标 修复前
P0 / P1 派发数 40 / 0
门闩触发次数 0

GSM8K 200q(高 cache 命中)也只能贴着 1.2× 边缘跑(1.197×),门闩 18/200 次,靠的是另一侧累计值恰好较低,并非设计正确。

修复

均衡门闩和 _select_worker_min_dispatched 兜底都改用按半衰期衰减的「近期」派发量 recent_dispatched_chars

  • 节点空闲时其近期值衰减回 0,门闩即能识别「曾经忙、现在闲」的倾斜并把流量转回去;
  • time.monotonic() 做 lazy 衰减,每次选点前衰减一次,无后台线程;
  • 默认半衰期 60s,可在 CacheAwareConfig.balance_half_life_secs 调;
  • 原累计值 dispatched_prompt_chars 保留(仅作统计),门闩已完全切到 recent。

涉及文件:

  • pd_io_struct.pyPD_Client_Obj 增加 recent_dispatched_chars / last_decay_ts
  • cache_aware.py:新增 _decay_recent,门闩与兜底改用 recent。
  • pd_selector.py:派发时累加 recent。
  • manager.py:节点注册重置时一并清零 recent。

验证

同样的 starvation 场景,修复后:

指标 修复前 修复后
P0 / P1 派发数(40 相同 prompt) 40 / 0 22 / 18
门闩触发次数 0 11

新增单测 unit_tests/server/test_cache_aware_balance.py(无需 GPU):

  • 衰减数学(两个半衰期 → 1/4);
  • 近期失衡时门闩重定向到空闲节点;
  • 负载均衡时仍保持 cache 亲和(回归保护)。

3/3 通过,black + flake8 通过。

说明

半衰期目前只在 CacheAwareConfig 可调,未接 CLI flag;如需要可后续加 --pd_balance_half_life_secs

…affinity starvation

The cache-aware prefill selector's 1.2x balance guard compared
cumulative-since-start dispatched chars, so a prefill node that was busy
early (now idle) kept a high cumulative value and the guard never fired
while cache affinity routed all new traffic to the other node, starving
the idle node. Repro: 40 identical prompts -> 40/0 split, guard fired
0 times.

Switch the guard (and the min-dispatched fallback) to a half-life-decayed
recent dispatched signal, so an idle node's load decays to 0 and the guard
redirects traffic back to it. Default half-life 60s, configurable via
CacheAwareConfig.balance_half_life_secs.

After the fix, the same workload splits 22/18 with the guard firing 11
times. Added unit tests for the decay math, the starvation redirect, and
that balanced load still keeps cache affinity.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant