Skip to content

feat(pd): add bounded decode-capacity admission control - #1529

Open
sufubao wants to merge 9 commits into
ModelTC:mainfrom
sufubao:feat/pd-cache-aware-admission
Open

feat(pd): add bounded decode-capacity admission control#1529
sufubao wants to merge 9 commits into
ModelTC:mainfrom
sufubao:feat/pd-cache-aware-admission

Conversation

@sufubao

@sufubao sufubao commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

背景与根因

main 中的 PD Master 在进行中请求数达到 Decode 容量时直接返回繁忙。本 PR 将它替换为有界、可取消的等待队列,并支持多 Master 容量切分、Session 优先级和 n choice 原子计费。

早期实现曾把 Prefill Radix cache 余量换算成冷请求并发硬上限:

cold_capacity = min(D, max(1, cache_free_tokens / avg_uncached_tokens))

Radix cache 满是可驱逐缓存的正常稳态,不是 Decode 并发安全边界。但该公式会在 cache 满时把冷请求并发压到 1;准入租约又持有到完整 stream 结束,最终使大量冷请求串行化。同时,节点每个 TOKEN_PACK 都扫描 Radix 共享状态,进一步放大了高压控制面开销。

最终方案删除这两条路径:Decode 容量是唯一的全局并发硬约束;Session 和缓存命中只影响软优先级。

准入模型

  • 汇总当前 Master 从所有 Decode 节点获得的容量份额,作为准入总容量。
  • 一个请求的 n 个 choice 原子占用 n 个 Decode 槽位;需求超过当前总容量时立即失败。
  • 租约在派发到 P/D 节点前获取,持续到正常结束、异常、取消或 stream close;获取后的设置异常也会统一释放。
  • 等待队列默认最多容纳一个 Decode 波次,按槽位而不是请求数计量。
  • 容量缩小时,不再可满足的 gang 请求立即失败;超过新队列上限的请求按低优先级、同级较新顺序移除。已发放租约不强制中断,在完成后收敛到新容量。
  • --disable_pd_master_decode_capacity_limit 继续用于完全绕过该准入队列。

公平调度与 gang backfill

请求入队时分为三类:

  1. CONTINUATION:服务端已观察到同一 X-Session-Id 的有效输出。
  2. PROBABLE_CACHE_HIT:cache-aware selector 估算的前缀命中率达到阈值。
  3. COLD:其余请求。

调度采用按 Decode 槽位计费的 deficit round-robin,默认 quantum 为 8:3:1;多 choice 请求按实际槽位成本扣费。

当一个已选 gang 只因当前空闲槽位不足而阻塞时,后续可容纳请求可以做有限 backfill;累计 backfill 达到一个 Decode 波次后转为 reservation,停止绕过该 gang,避免大 n 请求永久饥饿。

队列已满时,当前唯一可运行的新请求仍可在 backfill 预算内填充物理空槽;它不扩大等待队列,不越过实际可运行的 DRR 对手,也不破坏 Session FIFO 或 gang reservation。这避免了同 Session waiter 或多个 gang 占满队列时 Decode 槽位空转。

同一 Session 严格串行和 FIFO,其他 Session 仍可继续调度。排队请求支持按优先级超时和任务取消;队列满时,高优先级请求可以替换较低优先级等待项。

缓存信号与热路径清理

  • 缓存命中估算只影响准入软优先级,不再决定并发容量。
  • cache-aware 选点仍保留,并通过请求上下文复用准入阶段的 prefix match,避免同一请求重复遍历前缀树。
  • 删除 Radix cache 总 token、引用 token、容量、冷请求 EWMA 及相关 admission 状态。
  • TOKEN_PACK 和心跳只保留节点负载及 capacity_share / capacity_epoch
  • 只有有效 Decode 容量份额变化才重新驱动 admission;普通 token 负载、Prefill 上报和仅 epoch 变化不会重置 DRR/backfill 状态。
  • admission Gauge 为队列请求数、活动槽位和 Decode 容量;同一事件循环 tick 合并,相同值去重。

多 Master 与滚动升级

  • 每个 P/D 节点根据当前 Master 集合确定性切分 running_max_req_size;稳态时各 Master 份额之和等于节点容量。
  • Master 集合变化时推进 capacity_epoch 并立即唤醒心跳;旧 epoch 不会覆盖新份额。
  • 新 P/D 节点的 registration 保持旧 Master 可接受的顶层 schema,初始份额放在 start_args 保留键;旧节点连接新 Master 时仍回退到 running_max_req_size
  • 多 Master 滚动升级建议采用 P/D 节点优先,Master 随后 的顺序。如果先升级多个 Master 而仍有旧 D 节点,旧节点无法上报份额,各 Master 会保留旧版本的全容量回退行为。

验证

PYTHONPATH=. exp -m "final PR1529 bounded decode admission regression suite" \
  python -m pytest -q \
  unit_tests/server/test_pd_admission.py \
  unit_tests/server/test_pd_cache_aware.py \
  unit_tests/server/test_pd_master_mode.py \
  unit_tests/server/httpserver/test_pd_generate_error.py \
  unit_tests/server/httpserver/test_pd_master_cached_tokens.py \
  test/test_pd_selector/test_pd_master_multi_choice.py \
  test/test_api/test_server_busy_handling.py

104 项相关测试通过,覆盖:

  • Decode 容量边界、n choice 原子槽位、容量动态缩放和队列裁剪
  • 按槽位计费的 8:3:1 DRR
  • gang bounded backfill、满队列 idle-fill、reservation、取消和超时
  • Session 串行、FIFO、提升和队列替换
  • 多 Master 容量切分、epoch 和滚动升级协议兼容
  • Decode 租约的完整生成生命周期及获取后异常清理
  • 热路径遥测清理及 Gauge 合并去重

此外,提交前通过:

  • 200 seeds × 500 次 acquire/release/cancel/promote/capacity-change 随机状态转换不变量检查
  • Black 120、Flake8、git diff --checkpy_compile

控制器级合成实验(每请求模拟 50 ms 服务时间):

场景 中位耗时 有效速率 peak active
早期实现,D=64,cache full,64 请求 3.275 s 19.5 req/s 1
最终实现,D=64,64 请求 0.0509 s 1256.2 req/s 64
最终实现,D=64,128 请求 0.1026 s 1248.0 req/s 64

这是 admission controller 合成实验,不是真实模型/GPU 端到端吞吐数据。

范围与未验证边界

  • 本 PR 不修改 Radix cache 驱逐或分段策略,也不承诺通过 admission 控制缓存污染。
  • gang reservation 在一个 backfill 波次后可能有意保留空槽,这是防止大请求饥饿的公平性取舍。
  • Decode 租约覆盖完整请求生命周期;本 PR 不是独立建模 Prefill/Decode 两阶段容量的流水线调度器。
  • continuation 优先级依赖客户端显式发送 X-Session-Id;现有 benchmark_multiturn.py 尚未发送该请求头,因此没有端到端验证该优先级。
  • 尚未对最终实现运行真实模型的 GPU 高压吞吐、TTFT/P99 实验,也未实测 config-server 抖动下的多 Master 收敛过程。
@sufubao
sufubao force-pushed the feat/pd-cache-aware-admission branch from 93312ef to b915415 Compare August 31, 2026 12:11
@sufubao sufubao changed the title feat(pd): add cache-aware waiting queue admission Aug 31, 2026
@sufubao sufubao changed the title feat(pd): 添加缓存感知的等待队列准入策略 Aug 31, 2026
@sufubao sufubao changed the title feat(pd): add cache-aware waiting queue admission Aug 31, 2026
@sufubao sufubao changed the title feat(pd): add adaptive cache-aware admission control Aug 31, 2026
@sufubao

sufubao commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator Author

极压回退修正已推送:81b4cfa2

这次修正不再把 Radix cache 余量当作并发安全边界:

  • 删除 cache-full 时可退化到 1 的 cold hard cap,Decode 份额成为唯一全局并发硬约束。
  • 删除 TOKEN_PACK/心跳中的 Radix 共享状态扫描和无效 admission redrain。
  • 改为按 Decode slots 计费的 DRR + 一波 bounded backfill/reservation,并修复默认满队列时可运行小请求被 429、槽位空转的反例。
  • 修复新 P/D 节点连接旧 Master 的 registration schema 兼容性,并补齐 acquire 后异常的租约释放。

验证:104 项相关测试、200×500 随机状态转换以及 D=64 控制器饱和实验通过;GitHub pre-commit CI 已通过。控制器合成实验中,早期 cache-full 实现为 peak_active=1 / 19.5 req/s,最终实现为 peak_active=64 / 1256.2 req/s。这些不是真实模型/GPU 端到端吞吐数据,详细范围与未验证边界已写入 PR 正文。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant