Internal Wiki

主题

SGLang PD 分离实现速览 ​

双边对照的详细版本（sglang × tpu-inference）参见 sglang-tpu-inference-pd-disaggregation.md。

1. 为什么要 PD 分离 ​

prefill 与 decode 部署在同一个server时，由于优化目标不同，会在同一batch 里互相挤压；分开部署后两端各自优化对应的指标。

维度PrefillDecode
计算一次喂全部 prompt，矩阵乘大每步 1 token，矩阵乘小
资源compute-boundmemory-bound
指标TTFT（首 token 时延）ITL（token 间时延）
Batch拉长 step → TTFT 恶化摊薄读权重 → 收益高
  • 单server部署：TTFT和decode吞吐很难做到trade off
  • 分离：P 能够通过调整chunk_size，增加p节点数目来优化 TTFT，D可以通过攒更大的batch_size来获取更高的吞吐性能，独立扩缩容
  • 代价：KV cache 要显式从 P 传给 D，相比单server而言，有额外的传输开销

2. SGLang PD 全景 ​

5 类角色、3 类通信链路；进程粒度由启动参数决定。

  • MPMD：每 TP rank 一进程，KVManager per-rank
  • 角色靠 --disaggregation-mode prefill|decode 启动开关
  • P/D 的 TP 数可以不同（按 compute / memory 偏重差异化）
  • Bootstrap 只走元数据，KV 传输走 backend
  • Proxy 不在 sglang 主仓库，是部署细节

3. 一条请求的生命周期 ​

4 个阶段：路由 → P prefill+send → D pull → D decode。

  • Proxy 负责「P 完成了什么、放在哪儿」给 D
  • P 端发 KV 绕过 tree_cache，直读 req_to_token_pool
  • D 端预分配把 running / pre-allocated / transferring 三段隔离
  • 完成靠 KVReceiver.poll()==Success 信号

4. 控制面三件套 ​

4.1 Proxy 路由 ​

  • sglang 主仓库不带标准 Proxy，由部署方实现
  • 约定：先 max_tokens=1 的 prefill → 拿 kv_transfer_params → 再 stream=True 的 decode
  • 策略：round-robin / load-aware / sticky-to-room 自选
  • 参考：tpu-inference examples/disagg/toy_proxy_server.py

4.2 Scheduler Mixin ​

Mixin 而非继承，避免与 normal/overlap/pp/pdmux 组合爆炸。

  • 启动按 (disagg_mode, overlap, pp, pdmux) 选 event loop，如 event_loop_normal_disagg_prefill
  • P 端两队列：disagg_prefill_bootstrap_queue / disagg_prefill_inflight_queue
  • D 端两队列：disagg_decode_prealloc_queue / disagg_decode_transfer_queue
  • MPMD：每进程一份 Mixin；迁 SPMD 反而更简单

Scheduler 主循环 vs KV 传输的并行机制 ​

关键事实：KVSender.send 是异步入队，实际 RDMA 由 backend 自己的后台线程做——main loop 从不阻塞等传输。

  • normal 模式：每 tick 顺序 recv → pop_bootstrapped → run_batch → process_batch_result(含 send) → poll inflight_queue；send 异步，所以 batch N 的传输和 batch N+1 的 forward 自然 overlap
  • overlap 模式：用 self.result_queue = deque() 把 (batch, batch_result) 暂存一个 tick，下一 tick 才 process_batch_result——本 step 的 GPU forward 跑的同时 CPU 在处理上一步 sample + send，CPU 和 GPU 完全 pipeline
  • D 端对称：KVReceiver 的 pull 也在 backend 后台线程做（如 mooncake 的 start_decode_thread），main loop process_decode_queue 只做 prealloc + poll，KV 传输与 decode forward 全程 overlap
  • Scheduler ↔ Backend 无锁握手：KVPoll 状态由 backend 维护一张 check_status 哈希表，main loop 每 tick 读一次；backend 的握手/重试/超时全压在 poll() 后面，Scheduler 透明

4.3 KV 4-tuple ABC 与 backend ​

Base 类粒度职责
BaseKVManagerper-processbootstrap 注册、工厂
BaseKVSenderper-requestP 端发 KV
BaseKVReceiverper-requestD 端收 KV
BaseKVBootstrapServerper-cluster会合协议
  • 5 个 backend：Mooncake / NIXL / MoRI / Ascend / Fake，由 utils.py:get_kv_class 工厂分派
  • Sender/Receiver 用 per-request：每请求独立维持 uuid / buffer / 状态
  • backend 内部差异（RDMA、partial pull 等）对 Scheduler 透明，本文不展开

5. 多 pool 场景的差异封装 ​

Mamba / SWA / NSA 等额外 pool 都要随主 KV 传 P→D。sglang 用「接口扁平 + 字符串分派」吸收差异。

  • KVArgs 扁平化：3 组指针 + 1 个 state_type 字符串
  • 接口只增不变：send() 多挂可选 state_indices
  • 差异散落在每个 backend 的 maybe_send_extra 分支，没有公共 hybrid 基类
  • D 端：mamba 单独有 HybridMambaDecodeReqToTokenPool；SWA 复用普通子类，映射下沉到 allocator
  • 代价：新加一种 pool，每个 backend 都要重写一遍分支

6. 小结 ​

  • 三层解耦：路由层（Proxy）/ 控制层（Mixin + KVPoll）/ 数据层（4-tuple ABC + backend）
  • 想读代码：从 disaggregation/base/conn.py 开始 → 看 prefill.py / decode.py 的 Mixin
  • 双边对照：sglang-tpu-inference-pd-disaggregation.md