Skip to content

feat(pd): add node-authoritative Decode admission - #1529

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

feat(pd): add node-authoritative Decode admission#1529
sufubao wants to merge 10 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 容量时直接返回繁忙。此前版本把 admission 扩展成了 Master 侧的静态容量切分和业务调度器,但这会带来两个问题:

  • 多 Master 静态切份不是 work-conserving,空闲 Master 的份额无法被繁忙 Master 使用。
  • Decode lease 在 Master 进入生成流程前获取,覆盖多模态预加载、tokenization 和 Prefill,导致尚未进入 Decode 的请求提前占用槽位。

本版本改为:Decode 节点维护唯一、权威的本地容量池;Master 只负责在多 choice 请求完成 Prefill 后协调一次原子 n 槽预留。

实现

Decode 节点侧 admission

  • 每个 Decode 节点按本地 running_max_req_size 创建容量池,所有 Master 竞争同一个池。
  • 调度策略缩减为严格 FIFO,不包含优先级、DRR、Session 串行、驱逐、backfill 或调度 reservation。
  • 等待队列按槽位计量并有界,支持任务取消和超时。
  • n 个槽一次获取,不能部分准入。
  • lease 持续到对应 shm request index 真正回收,异常初始化和取消路径也会释放。
  • shm request index 增加原子批量分配:无法一次获得全部 index 时回滚,不留下部分占用。

Prefill / Decode 两阶段

  • n = 1 请求在 Decode 节点完成本地校验后、分配 shm request 前申请 lease。
  • n > 1 仍展开为独立的 n = 1 choice,以保持现有每个 choice 一条 KV 转移链路。
  • 所有 choice 完成 Prefill 后,Master 在同一个 Decode 节点原子预留 n 个槽;节点确认后才发送 Decode 请求。
  • 原子 lease 被拆成独立子 lease,各 choice 在自身物理资源回收时分别释放。
  • n > 1 不做 max_new_tokens 分段,避免不同 continuation history 被错误拼接;n = 1 保留原有分段行为。

多 Master 与滚动升级

  • 删除 capacity_share / capacity_epoch 静态切分和 Master 侧容量汇总。
  • Decode registration 上报节点侧 admission capability。
  • 新 Master 默认拒绝未提供该 capability 的旧 Decode 节点,要求按 Decode 优先、Master 随后 的顺序升级。
  • 显式禁用 admission 时允许连接旧 Decode 节点,便于紧急兼容。
  • 多机 TP 只由 rank 0 注册为 PD 入口并执行 admission;slave rank 只处理 rank 0 转发的请求。

错误与配置

  • Decode 队列满或等待超时时通过独立协议上报,最终保持 HTTP 429 ServerBusyError 语义。
  • 新增:
    • --pd_node_decode_admission_queue_size,默认等于 running_max_req_size
    • --pd_node_decode_admission_timeout,默认 5 秒。
    • --disable_pd_node_decode_admission
  • --disable_pd_master_decode_capacity_limit 保留为 deprecated alias。

精简结果

  • 删除原 Master admission controller 及其复杂策略测试。
  • 相对上一版 PR:1201 行新增、2310 行删除,净减少 1109 行。
  • cache-aware selector 继续负责选点,但不再参与 admission 优先级或容量计算。

验证

变更文件通过 Black 120、Flake8 和 git diff --check

python -m pytest -q \
  unit_tests/server \
  test/test_pd_selector \
  test/test_api/test_server_busy_handling.py

结果:310 passed, 1 failed。唯一失败为已有的:

unit_tests/server/router/dynamic_prompt/test_radix_cache.py::test_case10

该测试在 RadixCache.flush_cache() 中要求未初始化的 mem_manager 非空;对应实现和测试相对 upstream/main 均无改动。

本次未运行真实模型/GPU benchmark。

@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 feat(pd): 添加缓存感知的等待队列准入策略 Aug 31, 2026
@sufubao sufubao changed the title feat(pd): 添加缓存感知的等待队列准入策略 feat(pd): add cache-aware waiting queue admission Aug 31, 2026
@sufubao sufubao changed the title feat(pd): add cache-aware waiting queue admission feat(pd): add adaptive cache-aware admission control Aug 31, 2026
@sufubao sufubao changed the title feat(pd): add adaptive cache-aware admission control feat(pd): add bounded decode-capacity 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 正文。

@sufubao sufubao changed the title feat(pd): add bounded decode-capacity admission control feat(pd): add node-authoritative Decode admission Sep 1, 2026
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