RFC-2: PD 分离（Prefill-Decode Disaggregation）

字段	值
作者	@john
日期	2026-05-20
状态	Draft
目标	cross-process PD 分离：JaxTransfer ABC（参考 sglang 4-tuple 概念结构）+ Scheduler Mixin + `jax.experimental.transfer` backend；1P1D 端到端，MP+ND 接口预留
前置依赖	[[rfc_0_shared_cache_pd_infra]] M0（D 默认 ChunkCache）
关联 RFC	[[rfc_1_hicache]]（独立，同启时需协调 host memory）
关联 reference	[[sglang-tpu-inference-pd-disaggregation]]

1. 概述

协议层：JaxTransfer ABC，参考 sglang 4-tuple 概念结构（KVManager / KVSender / KVReceiver / BootstrapServer），形似但非签名兼容（§5.1）
Backend：JaxTransferKVEngine 承担 P↔D 跨进程 KV 传输；物理链路 D2H → H2H over TCP/gRPC → H2D（§5.2）；基于 jax.experimental.transfer + ZMQ 侧通道
Bootstrap：HTTP（完整 rank tuple + 多端口 schema，§5.3）
Proxy：直接复用 sglang sglang_router/mini_lb.py（python -m sglang_router.launch_router --pd-disaggregation，详见 §5.7 / §6.1）
Scheduler：Mixin 分隔 P/D event loop（§5.4）；D 端 pull_executor 包裹完整 conn.pull + wait_ready + insert_kv_chunks + ZMQ done 链路
D2H staging：第一版默认 off（§5.5）
部署：1P1D 端到端 + Proxy（P round-robin / D sticky-to-room）；MP+ND ABC 支持但不端到端验证

不实施：In-process 部署、mooncake/NIXL/MoRI、partial KV pulling、PP+PD、Spec Dec+PD、SWA+PD、Mamba+PD。

2. 背景与现状 gap

2.1 现状 vs sglang / tpu-inference 对比

维度	sgl-jax 现状	sglang (GPU)	tpu-inference (TPU PoC)	sgl-jax 第一版
协议层	无	4-tuple ABC	自定义 `TPUConnector*`	JaxTransfer ABC（参考概念结构）
Backend	无	mooncake / NIXL / MoRI / ZMQ	`jax.experimental.transfer`	`jax.experimental.transfer`
Bootstrap	无	HTTP + heartbeat	TODO	HTTP + heartbeat + TTL
D2H staging	无	mooncake 自动 stage	`TPU_ENABLE_D2H_TRANSFER` 默认 on	默认 off
D cache	无 PD	ChunkCache + opt-in RadixCache	复用 tpu-inference cache	RFC-0 已定 ChunkCache
Scheduler	单 event loop	P/D 各 mixin	—	P/D 各 mixin（移植）

详见 [[sglang-tpu-inference-pd-disaggregation]] §2/§4。

2.2 JAX backend 约束

约束	后果
`jax.experimental.transfer` 实验性	锁 JAX 版本；ABC 隔离便于切换
不支持 partial pull	一次拉整 shard
无 native async pull	ZMQ 侧通道做 ready 通知 + pull_executor 线程池
无 GPUDirect-RDMA 等价物	跨 process 物理必经 `D2H → H2H(TCP) → H2D`，无 D2D 直传

2.3 部署形态选择见 §4

In-process / local-disagg 是否实施、SPMD vs MPMD trade-off 详见 §4。

3. 架构总览

关键关系：

4-tuple ABC 是协议层抽象，KVManager 持 backend + per-request KVSender/KVReceiver
Bootstrap 是控制平面（实例注册 / 查询 / heartbeat），不参与 KV 数据传输
Proxy 直接复用 sglang sglang_router/mini_lb.py（asyncio.gather 同时 post P/D + 顶层 bootstrap_host/port/room 注入），sgl-jax 不写自有 Proxy；第一版 P round-robin / D sticky-to-bootstrap_room（§5.7 / §6.1）
JaxTransferKVEngine 是数据平面，进程内单例
ZMQ 侧通道 用于 ready / done 通知（绕过 jax.experimental.transfer 缺 async pull 的限制）

4. 部署形态：SPMD vs MPMD

PD 分离在 TPU/JAX 上有两种部署形态。本节梳理两者的物理 / 软件层约束、trade-off 与 sgl-jax 选择。

4.1 概念

形态	全称	JAX cluster 划分	P↔D 通信
SPMD	Single Program Multiple Data（In-process / local-disagg）	单 `jax.distributed.initialize` cluster + 单 SPMD program；P/D 通过 mesh 角色分区	XLA 编译期生成集合通信原语（`all_gather` / `ppermute` / `jax.device_put`），slice 内走 ICI
MPMD（第一版选择）	Multiple Program Multiple Data（cross-process）	P 与 D 各自独立 `jax.distributed.initialize` cluster，独立 PJRT device registry	4-tuple ABC + `jax.experimental.transfer`，物理 `D2H → H2H(TCP) → H2D`

4.2 物理 / 软件层约束（为什么 MPMD 即便同 slice 也走 host）

物理上 v5p-256 = 4×4×8 chip topology（参考 GCP v5p 文档；不同规格 slice topology 不同——如 v5p-128 = 4×4×4、v5p-512 = 4×4×16），同 slice 内 P/D 之间 ICI 链路存在；但软件层把 ICI 使用权"圈"在 cluster 边界内：

约束	影响
ICI 不跨 slice	即便 SPMD 单 program，跨 slice 也走 DCN（D2H→网络→H2D）；ICI 只在 slice 内 3D torus 物理可达（Google Cloud Multislice 文档）
`jax.distributed.initialize` 每进程仅一次	一个 process 只能属一个 cluster；JAX 架构禁止跨 cluster device 共享
PJRT device registry 绑定 cluster	独立 cluster 的 `PjRtDevice*` 在对方 process 根本没有对象——跨 cluster `CopyToDevice` 失败
libtpu 未暴露 cross-process IPC API	TPU 无 GPU CUDA IPC（`cudaIpcGetMemHandle`/`cudaIpcOpenMemHandle`）等价物。这是 MPMD 即便同 slice 内物理 ICI 可达，软件层也无法 D2D 的最终根因——`jax.experimental.transfer` 的 TCP-only 不是工程妥协，是 docstring 明示的"DCN cross slice transfer"设计目标 + libtpu 不暴露 IPC 的双重结果

4.3 Trade-off 与 sgl-jax 选择

维度	SPMD（In-process）	MPMD（cross-process，第一版）
slice 内 P↔D 传输	✅ ICI 直传（D2D zero-copy）	❌ `D2H → H2H(TCP) → H2D`（即便单机 localhost）
跨 slice / 跨机 P↔D	走 DCN，与 MPMD 等价	同
角色弹性扩缩	静态——P/D 比例编译期固定	动态——`disaggregation_mode` 决定本 process 角色，可独立扩缩
异构硬件	必须同 cluster 同硬件	P/D 可跨代次 / 跨机
单点故障	一 process 挂 P/D 同时丢	P/D 独立可独立重启
运维复杂度	低（单 process）	高（多 process + Bootstrap + Proxy）
代码复杂度	高（基础设施：多 sub-engine 的 device 分配 + 多 JetThread × 3 阶段 pipeline + queue-based KV 搬运 + load balance，参考 tpu-inference `DisaggExecutor` + `_DisaggOrchestrator`）；P/D forward 在各自 thread 内独立 dispatch 解耦	低（process 隔离自然 + 4-tuple ABC + scheduler mixin 解耦，无需多 thread 协调）
生态对齐	tpu-inference Local Disagg	sglang / vLLM PD 主流
GPU 对照	vLLM 单 process split-PD（少见）	sglang / vLLM 主流多 process PD

sgl-jax 第一版选 MPMD cross-process，理由：

生态对齐：sglang / vLLM 主流是 MPMD；跟随主流最大化复用 sglang_router/mini_lb、4-tuple ABC、scheduler mixin
弹性扩缩 + 异构部署：生产场景 P/D 的 compute/memory profile 不同（P 重 compute、D 重 KV memory），独立扩缩是核心需求
代码解耦：scheduler / KV manager / cache 在两种角色下完全独立，便于演化
跨 slice 性能等价：实际部署常跨机，SPMD 在跨 slice 时也只能走 DCN，ICI 优势消失

代价：放弃单 slice 内 ICI 直传的性能上界——但只在"单 slice 部署"场景才能拿到，跨机场景不适用。

4.4 未来 SPMD 适用场景（§9 范围外）

某些场景下 SPMD 仍有价值，作为未来工作明确条件：

单 slice 内固定 P/D 比例的小规模部署（如 v5p-256 内固定 2P+2D，不需弹性扩缩）
延迟敏感且不需弹性的内部 benchmark / 性能上界对比
跨 slice / 多机部署 SPMD 也救不了——物理 ICI 限制相同，不在 SPMD 适用范围

实施路径：参考 tpu-inference examples/disagg/single_proc_per_worker.py 的 Local Disagg 设计。

4.5 Cross-process 部署细节（第一版实施）

角色	进程数	模型
Bootstrap	1	HTTP 服务，不占 TPU
Prefill (P)	≥ 1	每实例 1 process，独立 TPU mesh
Decode (D)	≥ 1	同上
Proxy	1	sglang `sglang_router/mini_lb`（复用），不占 TPU

单机最小部署（4 process 共享 TPU，每 process 用 TPU_VISIBLE_CHIPS 静态切 chip）：

进程	TPU device	端口	JAX init
Bootstrap	不占	HTTP 8000	不调 `jax.distributed.initialize`
Proxy	不占	HTTP 8001	不调
P	`TPU_VISIBLE_CHIPS=0,1,2,3`（v6e-8 切一半举例）	transfer 7080 / ZMQ 9080 / IPC 6080	`jax.distributed.initialize(...)` 独立 mesh
D	`TPU_VISIBLE_CHIPS=4,5,6,7`	transfer 7081 / ZMQ 9081 / IPC 6081	同上

chip 划分由 launch script 静态指定，端口按"每角色 +1"；D3 阶段交付 launch_single_node_pd.sh 模板。多机部署同模式跨机切分；启动器（Ray / K8s / 手动）留实施层（§10）。

5. 模块设计

5.1 4-tuple ABC

参考 sglang 4-tuple 概念结构，sgl-jax 专属 JaxTransfer ABC——形似但非签名兼容：sglang BaseKVReceiver.send_metadata / clear / abort 是 mooncake/NIXL 数据面需要（先发 indices 再 RDMA 直传），JaxTransfer 走 await_pull / pull 不需要；不复用 sglang proxy/router/backend adapter；只保留概念分层与 KVPoll 状态机语义。

ABC	关键方法（JaxTransfer 专属）	实现
`BaseKVManager`	`__init__(args, bootstrap_addr)` / `register_runner(runner)` / `create_sender(room)` / `create_receiver(room)`	`JaxTransferKVManager`
`BaseKVSender`	`send(kv_indices)` / `poll() → KVPoll` / `failure_exception()`	内嵌
`BaseKVReceiver`	`__init__(mgr, bootstrap_addr, bootstrap_room)`（从 req 顶层字段 `bootstrap_host`/`port`/`room` 拿，与 sglang `base/conn.py` 签名一致；额外 KV meta 走 `kv_transfer_params` 子字段，详见 §6.2）/ `poll() → KVPoll` / `failure_exception()`	内嵌
`BaseKVBootstrapServer`	`start()` / `stop()` / `poll() → KVPoll`	`HTTPBootstrapServer`

KVPoll 状态机（与 sglang 一致）：Bootstrapping → WaitingForInput → Transferring → Success / Failed

register_runner 改造：scheduler 创建完 runner 后显式注入（kv_manager.register_runner(runner)），KVCache 不持反向引用，避免循环依赖。

5.2 JaxTransferKVEngine 数据链路

进程内单例 transfer_server（与 tpu-inference 一致）。物理链路始终 D2H → H2H over TCP/gRPC → H2D——无 D2D 直传选项，NIC 在 host CPU、TPU 无 GPUDirect-RDMA；即便单机 localhost 也走 TCP。

ZMQ 侧通道：

用途	流向	内容
ready 通知	P → D	`bootstrap_room`，shard meta
done 通知	D → P	`bootstrap_room`，success/fail

按 §5.5 D2H staging 开关切 P 端 D2H 由谁做（不是切是否走 host）：

模式	P 端 D2H	HBM 释放	第一版
A：disabled（默认）	transfer_server C++ runtime 在 D pull 时隐式	等 D pull 完	实施
B：enabled	P 端 Pallas `copy_to_host` 显式写预分配 `HostKVPool` pinned buffer	D2H 完成即可	不实施（单机性能优化，非目标场景；未来参考 tpu-inference `tpu_connector.py:712-760` + `kv_transfer.py:416`）

数据形态（paged ↔ packed ↔ paged）：device 端 paged pool 是 [total_blocks, page_size, ...] 大 buffer，单个请求的 block 在 pool 内不连续（由 page allocator 分散分配）。PD 传输全链路：

P 端 paged → packed gather（§5.2.1 P1）：按 local_block_ids 从 paged HBM 取出形状 [num_blocks, page_size, ...] 的紧凑 array list（每 layer 一个）；这是 jax.lax indexed gather，不是 view 而是拷贝
跨进程 packed 传输：await_pull / pull 注册和拉取的就是这个紧凑 list，每 layer 一个 [num_blocks, ...] array（D2H + 网络 + H2D 全程保持 packed 形态）
D 端 packed → paged scatter（§5.2.1 D2d）：按 D scheduler preallocate 的 block_numbers（也可能不连续），用 multi_layer_copy（Pallas DMA + 跨 layer 流水线）把紧凑 list 分段 scatter 回 D 端 paged pool 的指定 block 位置

5.2.1 P 端 `KVSender.send`（模式 A）

#	调用	物理
P1	`kv = select_from_kv_caches(runner.kv_caches, indices)`	paged → packed gather：`cache.at[indices].get()`（`jax.lax` indexed gather，不是 view 是拷贝），从 paged HBM `[total_blocks, ...]` 取出形状 `[num_blocks, page_size, ...]` 紧凑 array list（每 layer 一个），参考 `tpu_connector.py:914-917`
P2	`reqs_wait_pull[req_id] = (kv, expiration, -1)` + `uuid → req_id` 映射	留 HBM 引用 + 记账
P3	`transfer_server.await_pull(uuid, kv)`	非阻塞注册
P4	ZMQ notify D：`uuid` ready	侧通道 push

收到 D3 done → 删 reqs_wait_pull[req_id] → HBM 引用归零自动释放。

5.2.2 D 端 `KVReceiver.pull`（模式 A）

conn.pull 是阻塞调用 + 返回未必 ready 的临时独立 buffer（不在本地 paged pool）——必须由 pull_executor: ThreadPoolExecutor 调度，且在后台做 readiness wait + scatter；主循环查 future.done()。参考 tpu-inference tpu_connector.py:666-668 insert_kv_chunks + :797-805 is_ready() true 分支。

#	调用方 / 调用	物理
D1	主线程：`conn = transfer_server.connect(p_addr)`（或复用）	跨 process 连接
D2a	主线程：`future = pull_executor.submit(_pull_kv, req_id, conn, req_meta)`	提交后台；主线程立即返回
D2b	后台线程：`kv = conn.pull(uuid, kv_spec)`	TCP 拉取；返回 future array，未必 ready，临时独立 buffer
D2c	后台线程：readiness wait——第一版永远 wait（不论环境变量），轮询 `chunk.is_ready() for chunk in kv` + 超时；不学 tpu-inference fire-and-forget false 分支	阻塞到 ready
D2d	后台线程：`runner.kv_caches = insert_kv_chunks(runner.kv_caches, kv, block_numbers, mesh, sharding.spec)`	packed → paged scatter：内部按 `block_numbers` 计算连续段（pad 到 total_len 避免重编译），调 `multi_layer_copy`（Pallas DMA + 跨 layer 流水线，`kv_transfer.py:296`）把临时紧凑 kv list 分段 scatter 回本地 paged pool；之后 D forward 才能从 pool 读到正确 KV
D3	后台线程：ZMQ notify P：done；future 标记 done	侧通道 push

future.done() = True 严格等价"D2b + D2c + D2d 全部完成、本地 paged pool 已写好"。 主循环 process_decode_queue 见 done 即把 req 入 waiting_queue 进 decode。

5.2.3 其他约定

不支持 partial pull（API 限制）：D2b 一次拉完 kv_spec 指定 shard
pull 落地由 kv_spec.sharding 决定：第一版统一 device HBM sharding
失败处理：transfer 失败 → KVPoll = Failed → scheduler 上报客户端
buffer 生命周期：P 端 expiration_time 超时清理孤儿（D 不来 pull 的情况）

5.3 BootstrapServer（HTTP）

维护 P/D 实例注册表（完整 rank tuple + 多端口），支撑 MP+ND。不维护 room→P 配对——由 Proxy 在 req 顶层注入 bootstrap_host / bootstrap_port / bootstrap_room 三字段（与 sglang mini_lb.py:358-405 + sglang base/conn.py:113-121 顶层字段约定一致；额外 KV meta 入 kv_transfer_params 子字段，详见 §6.2）。

端点	用途
`PUT /route`	实例注册（payload 见下）
`GET /list?role=prefill` 或 `?role=decode`	Proxy 枚举做 round-robin / sticky
`GET /route?role=...&instance_id=...&tp_rank=N&cp_rank=N&dp_rank=N&pp_rank=N`	TP/DP/PP 多 rank 协调按完整 rank tuple 查
`POST /heartbeat`	默认 5s 间隔
`GET /health`	Liveness

PUT /route payload schema：

text

{
  role: "prefill" | "decode",
  instance_id: str,                          # 区分多实例
  tp_rank, tp_size, cp_rank, cp_size,
  dp_rank, dp_size, pp_rank, pp_size: int,   # 完整 rank tuple，避免多实例 key 冲突
  transfer_addr: "ip:port",                  # jax.experimental.transfer server
  zmq_addr: "ip:port",                       # ZMQ 侧通道
  scheduler_addr: "ip:port",                 # scheduler IPC（可选）
  page_size, kv_dtype, mesh_meta: ...
}

key = (role, instance_id, tp_rank, cp_rank, dp_rank, pp_rank)。bootstrap_room 64-bit 整数由 Proxy atomic counter / UUIDv4 截断生成，不做服务端反查校验。

TTL：heartbeat_ttl = 15s（= 3 × interval），超时移除不主动通知（对端通过 ZMQ 失败感知）。

第一版限制：单点无 HA；无持久化（重启后实例重新注册）；端到端验证 1P1D，MP+ND schema 就位但不端到端验证。

端点风格参考 sglang common/conn.py:754-799（完整 schema 含 attn_tp_size/attn_tp_rank/... 等字段）。

5.4 Scheduler Mixin

Mixin	何时挂	提供
`SchedulerDisaggregationPrefillMixin`	`disaggregation_mode=prefill`	`_event_loop_prefill_disagg` / `_create_send_queue` / `_handle_completed_prefill`
`SchedulerDisaggregationDecodeMixin`	`disaggregation_mode=decode`	`_event_loop_decode_disagg` / `_create_recv_queue` / `_handle_recv_complete`

装配：

text

class PrefillScheduler(SchedulerDisaggregationPrefillMixin, Scheduler): ...
class DecodeScheduler(SchedulerDisaggregationDecodeMixin, Scheduler): ...

def make_scheduler(args):
    if args.disaggregation_mode == "prefill": return PrefillScheduler(args)
    elif args.disaggregation_mode == "decode": return DecodeScheduler(args)
    else: return Scheduler(args)

Event loop 主路径：

阶段	主路径
P	recv req → forward → `KVSender.send`（`await_pull` 非阻塞注册 + ZMQ notify D）→ 释放 req
D	读 decode req 顶层 `bootstrap_host`/`port`/`room` + `kv_transfer_params` 额外 meta → `KVReceiver.__init__(mgr, bootstrap_addr, bootstrap_room)` → `pull_executor.submit(_pull_kv)`（封装 `conn.pull → wait_ready → insert_kv_chunks → ZMQ done notify`，§5.2.2）→ `future.done()` → req 入 `waiting_queue` → forward decode → EOS

并行机制：scheduler 主循环单线程，每 step 非阻塞 poll：

端	关键调用	阻塞性	调度
P	`transfer_server.await_pull(uuid, kv)`	非阻塞	主线程裸调，进 `reqs_wait_pull`
D	`conn.pull(uuid, kv_spec)`	阻塞	必须 `pull_executor: ThreadPoolExecutor` 提交（参考 tpu-inference `tpu_connector.py:514, 653`）

主循环只做"提交 + poll + 搬运"——forward 在 device 异步执行（JAX dispatch），transfer 在 C++ runtime + pull_executor 后台并行；不阻塞主循环。

D 端时序：

P 端时序类似但更简：P 端不需要 pull_executor，主线程裸调 await_pull 非阻塞注册；只多一个 ZMQ 监听线程把 D 的 done 通知写共享状态。

Thread 总览：

角色	用途
主循环线程	event loop（提交 / poll / 搬运），P/D 共用
`pull_executor`	D 端必需，包裹 `conn.pull + wait_ready + insert_kv_chunks + ZMQ done` 完整链路
ZMQ 监听线程	收对端 ready / done 通知
`kv_d2h_executor`	§5.5 模式 B 才需要，P 端可选

Overlap 模式（event_loop_overlap_disagg_*）再叠 step 间 forward 与 process_batch_result 的 overlap（result_queue = deque() 把当前 step result 延后一个 step），仍单线程。

与 tree_cache 交互：PD 主路径绕过 tree_cache；D 默认 ChunkCache（RFC-0 已定），lookup 必 miss，不浪费 tree 操作。

5.5 D2H staging（第一版 off）

切 P 端 D2H 由谁做（不切是否走 host）；两种模式物理链路都是 D2H → H2H over TCP → H2D：

模式	P 端 D2H	HBM 释放	适用
`disabled`（默认）	transfer_server C++ runtime 隐式（可能触发 `MapDmaBuffer` 动态映射）	等 D pull 完	跨机 / 小 shard / 第一版统一
`enabled`	P 端 Pallas `copy_to_host` 显式写 `HostKVPool` 预分配 pinned buffer	D2H 完成即可	单机 / 大 shard / HBM 紧

TPU_ENABLE_D2H_TRANSFER 默认 0（off）——与 tpu-inference 上游（默认 on）相反，sgl-jax 第一版统一 off 以最小化工程量；D3 之后按 profile 决定是否开。详见 [[sglang-tpu-inference-pd-disaggregation]] §4.3。

5.6 接口预留：多 pool transfer descriptor

第一版只传 main KV pool；未来扩展时单请求要传 N 个 pool（Mamba/KDA 的 RecurrentStatePool、Spec Dec 的 draft KV pool、SWA 的 valid range metadata）。一个 enum 不够——sender/receiver 必须能同时传"主 + sidecar"复合形态。

ABC 用 KVArgs.pools: List[PoolTransferSpec]（每项 {name, pool_type, dtype, shape, sharding, optional metadata}）。sender/receiver init 时遍历 pools 各自 await_pull / pull；ZMQ notify 一次性触发全部 pool 的 transfer，KVPoll 在全部完成后置 Success。

模型类	`pools` 配置
MHA / GQA / MLA + standard	`[main_kv]`（第一版唯一）
Mamba / KDA	`[main_kv, recurrent_state]`
Spec Dec（MTP / EAGLE，sgl-jax PR #1108）	`[main_kv, draft_kv]`
SWA	`[main_kv]` + `KVArgs.attention_window` metadata

sglang KVArgs.state_type 是 enum 风格（none / mamba / swa），sgl-jax 选 list 是为显式支持"主 + sidecar"；因 JaxTransfer ABC 本就不承诺与 sglang 签名兼容（§5.1），enum→list 不增加协议偏离成本。

第一版：实装单 main_kv；ABC 层接受 list；多 pool 端到端归未来工作。

5.7 MP+ND

多 P + 多 D 并行部署，Proxy 在请求级别选实例。Proxy 复用 sglang sglang_router/mini_lb.py（启动：python -m sglang_router.launch_router --pd-disaggregation --prefill <P_url> --decode <D_url>），sgl-jax 不写自有 Proxy。BootstrapServer 注册实例 list；bootstrap_room 1:1 配对路由，不与实例数耦合；KVManager 通过 req 顶层 bootstrap_host / bootstrap_port 拿 P transfer_addr，跨 rank 协调按 rank tuple 查 GET /route。

第一版：

1P1D 端到端跑通（功能正确性 + 性能 baseline）
Proxy 路由：P round-robin（每新请求轮询）+ D sticky-to-bootstrap_room（同 room 始终路由同 D）
2P+2D 接口能跑（Proxy 轮询能工作），但不做性能调优、不做智能路由

智能调度（prefix-aware / load-aware / cache-aware）归 §9 范围外。多机启动器（Ray / K8s / launch script）留实施层（§10），D3 阶段提供 launch script 模板。

6. 关键数据流

6.1 P → D KV 传输（模式 A，默认）

模式 B 端到端时序相同，只把 P 侧隐式 D2H 替换为 §5.2/§5.5 描述的显式 Pallas copy_to_host + 预分配 HostKVPool 路径。

关键约定：

Proxy 同时路由给 P 与 D（asyncio.gather），req 顶层注入 bootstrap_host / bootstrap_port / bootstrap_room 三字段（与 sglang mini_lb.py:358-405 完全一致）；额外 KV meta（shape/sharding/page_size）入 kv_transfer_params 子字段（详见 §6.2）
与 sglang 官方 sglang_router/mini_lb.py:135-141 一致：mini_lb 用 asyncio.gather 同时 post P+D；_generate_bootstrap_room() 由 mini_lb 生成；P/D 从 req 顶层字段读 bootstrap 信息（tokenizer_manager.py:999-1001 透传）。sgl-jax 直接复用 sglang_router/mini_lb，不写自有 Proxy 代码；只提供 launch script
这一模式偏离 tpu-inference toy_proxy_server.py:201 的串行流程（toy proxy 是 P 完成 → callback 拿 kv_transfer_params 回流 → 再发 D），sgl-jax 跟随 sglang 主流做法
D 端 conn.pull 阻塞，必须由 pull_executor 调度；主循环每 step 查 future.done()
物理上无 D2D 直传（§5.2）

6.2 Proxy 注入字段 schema（顶层 + 子字段两层）

顶层字段（sgl-jax 与 sglang mini_lb.py:369-381 一致，由 mini_lb 同时塞入 prefill + decode req）：

字段	来源	用途
`bootstrap_host`	mini_lb 从 `--prefill <url>` 解析的 P entry hostname（multi-host P 时一般是 host 0）	D 端获取入口；具体连哪个 P host 见下文"multi-host P N:N 配对"
`bootstrap_port`	mini_lb 从 `select_pair()` 拿 P 的 bootstrap_port	同上
`bootstrap_room`	mini_lb `_generate_bootstrap_room()`（64-bit）	作 `uuid` 在 `await_pull` / `pull` / ZMQ 间贯穿
`routed_dp_rank` / `disagg_prefill_dp_rank`	mini_lb `_fork_dp_requests()`（仅 dp_size>1）	DP 路由

kv_transfer_params 子字段（sgl-jax 在 sglang 基础上扩展，承载 D 调 pull 所需的额外 metadata）：

字段	来源	用途
`zmq_addr`	Bootstrap `GET /list` 拿 P 实例 `zmq_addr`	D 端 ZMQ ready 通道
`kv_shape` / `kv_dtype` / `sharding_meta` / `page_size`	同上，P 实例模型 meta	D 端构造 `kv_spec` 调 `pull`
`pool_specs`	对应 §5.6 `KVArgs.pools`	第一版 `[main_kv]`

block_numbers 不由 Proxy 注入——由 D scheduler preallocate 阶段产生，在 insert_kv_chunks 时从 req_meta.local_block_ids 拿。
expiration_time 不由 Proxy 注入——P 端在 _prepare_kv_and_wait 本地计算（now + timeout），用于清理孤儿 buffer（§5.2.3 buffer 生命周期）。

Multi-host P 的 N:N 对称配对模式：

jax.experimental.transfer 接口本身只提供 per-process 原语（start_transfer_server 单 host listen + TransferConnection 是 1:1 peer，参考 transfer.py:21 docstring "Represents a connection to exactly one peer"），多 host 分布式传输由调用层显式拼接（不是接口自动 multi-source）。tpu-inference 的模式（tpu_connector.py:763-765）：

每个 P host 各自启 transfer_server（listen 自己的 host_ip:port），注册本 host 持有的 KV shard
每个 D host 用自己的 node_id 索引 P host 列表，连对应索引的那个 P host——remote_addr = f"{remote_host[self.node_id]}:{remote_port[self.node_id]}"
这是 N:N 对称配对：P host i ↔ D host i，前提是 P/D 节点数 + sharding layout 一致（同一份 mesh 在两侧复用，这是 LLM serving 的常见假设）

D 端获取对应 P host addr 的实际路径：D 端 KVReceiver init 时拿到顶层 bootstrap_host 作为入口，按自己的 (tp_rank, cp_rank, dp_rank, pp_rank) 调 Bootstrap GET /route?role=prefill&instance_id=...&tp_rank=N&cp_rank=N&dp_rank=N&pp_rank=N 查对应 P host 的真实 transfer_addr（每个 P host 在启动时已 PUT /route 注册，§5.3 schema 已含 transfer_addr 字段）。这样 mini_lb 注入的顶层 bootstrap_host 单值与 sglang schema 兼容，multi-host N:N 配对由 D 端按需查 Bootstrap 完成。

为什么顶层 + 子字段两层：顶层 bootstrap_* 复用 sglang mini_lb 现成 schema，sgl-jax 不需要 patch mini_lb；kv_transfer_params 子字段是 sglang 已有的概念（tokenizer_manager.py 内部也用），承载 sglang 没有的 JAX backend 专属 meta（shape/sharding/pool_specs 等），与 mini_lb 解耦。

Bootstrap 注册 / heartbeat / 失联流程见 §5.3。

6.3 PD + ChunkCache（D 视角）

D scheduler → RFC-0 builder → cache_type=chunk_cache → ChunkCache 实例。新 req → tree_cache.match_prefix 必 miss → 走 PD 主路径（pull_executor 完成 conn.pull + wait_ready + insert_kv_chunks，本地 paged pool 写好）→ forward decode。与 sglang ChunkCache + PD 路径一致。

7. 配置项

text

--disaggregation-mode {null, prefill, decode}
    null  → 不启 PD
    prefill / decode → 本进程作为 P/D（D 的 cache_type 默认值在 RFC-0 已定为 chunk_cache）

--disaggregation-transfer-backend {jax}              # 第一版唯一值
--disaggregation-bootstrap-host HOST
--disaggregation-bootstrap-port PORT
--disaggregation-bootstrap-heartbeat-interval SEC    # 默认 5
--disaggregation-bootstrap-heartbeat-ttl SEC         # 默认 15（= 3 × interval）

TPU_ENABLE_D2H_TRANSFER={0,1}    # 第一版默认 0 (off); tpu-inference 上游默认 on
TPU_ENABLE_BLOCK_KV_TRANSFER     # 第一版不暴露——pull_executor 永远 wait before scatter,
                                 # 不学 tpu-inference false 分支 fire-and-forget (会让 D
                                 # forward 拿到未 ready array)

不引入：disaggregation_mode=local、其他 backend 选择、智能路由策略开关。

8. 实施路线

阶段	交付	验收	依赖
D0	`JaxTransferKVEngine` + `BootstrapServer` 单测	ZMQ + transfer 拷贝端到端单测；Bootstrap 三端点单测；不接入 scheduler	RFC-0 M0
D1	4-tuple ABC + Scheduler Mixin（PrefillScheduler / DecodeScheduler）+ ChunkCache 集成 + 接入 sglang `sglang_router/mini_lb`（启动 `python -m sglang_router.launch_router --pd-disaggregation --prefill <P_url> --decode <D_url>`）+ 单机 4-process launch script	1P1D 单机端到端 serve 小模型（via mini_lb，参考 sglang `disaggregation_fixture.py:104-132` 测试模式）；非 PD 路径 KL 等价；`register_runner` 显式注入；req 顶层 `bootstrap_host`/`port`/`room` + `kv_transfer_params` 子字段全链路追踪	D0
D2	UnifiedRadixCache + PD 兼容验证 + 跑 Llama/Qwen 真模型	1P1D + mini_lb 端到端跑 Llama/Qwen；P + UnifiedRadixCache + PD 三者兼容；性能 baseline	D1，[[rfc_1_hicache]] H1
D3	多机部署文档 + PD + HiCache 同启 benchmark	多机 1P1D 启动；HBM/host memory 切分合理；性能 baseline（vs 非 PD）	D2，[[rfc_1_hicache]] H3

与 RFC-1 协调：D2 要 RFC-1 H1 完成；D3 与 RFC-1 H3 同期 benchmark。

9. 范围外 / 未来工作

项目	原因
In-process 部署	与生产场景不符
Mooncake / NIXL / MoRI	GPU 专用
Partial KV pulling	API 不支持
PP+PD / Spec Dec+PD / SWA+PD / Mamba+PD	sgl-jax 暂未支持或 ABC 仅预留
跨进程共享 host pool	sglang 也未实现
Ray / K8s 启动器	RFC 不规定具体启动器
HA Bootstrap	短期接受单点
MP+ND 性能调优 + 智能路由	接口支持但第一版不调
D + RadixCache opt-in flag	sgl-jax 第一版不引入
D2H staging 默认启用	D3 后 profile 决定
`DecodeKVCacheOffloadManager`（D + L3）	RFC-1 + RFC-2 同启时再议
异构 sharding PD（P/D TP/DP/PP 数不同，如 P TP=8 + D TP=4）	第一版假设 P/D 节点数 + sharding layout 对称（N:N 配对，§6.2）；异构需要应用层 KV reshape adapter（split/merge shard），不在 `jax.experimental.transfer` 接口能力范围

10. 未决问题

问题	实施时定
多机启动器（Ray / K8s / launch script）	D3
PD + HiCache 同启 HBM/host memory 切分比例	D3 + H3 联合 benchmark
ZMQ 端口范围与冲突处理	D0
`bootstrap_room` 重复检测（64-bit 碰撞概率极低）	D0，倾向不引入
Proxy streaming 响应聚合	D2
失败请求客户端通知协议	D2
Heartbeat interval / TTL 默认值	D3 benchmark 后定

文档设计

tpu-inference Q2 对照

sglang 代码分析

RFC-2: PD 分离（Prefill-Decode Disaggregation）

1. 概述

2. 背景与现状 gap

2.1 现状 vs sglang / tpu-inference 对比

2.2 JAX backend 约束

2.3 部署形态选择见 §4

3. 架构总览

4. 部署形态：SPMD vs MPMD

4.1 概念

4.2 物理 / 软件层约束（为什么 MPMD 即便同 slice 也走 host）

4.3 Trade-off 与 sgl-jax 选择

4.4 未来 SPMD 适用场景（§9 范围外）

4.5 Cross-process 部署细节（第一版实施）

5. 模块设计

5.1 4-tuple ABC

5.2 JaxTransferKVEngine 数据链路

5.2.1 P 端 `KVSender.send`（模式 A）

5.2.2 D 端 `KVReceiver.pull`（模式 A）

5.2.3 其他约定

5.3 BootstrapServer（HTTP）

5.4 Scheduler Mixin

5.5 D2H staging（第一版 off）

5.6 接口预留：多 pool transfer descriptor

5.7 MP+ND

6. 关键数据流

6.1 P → D KV 传输（模式 A，默认）

6.2 Proxy 注入字段 schema（顶层 + 子字段两层）

6.3 PD + ChunkCache（D 视角）

7. 配置项

8. 实施路线

9. 范围外 / 未来工作

10. 未决问题

tpu-inference Q2 对照

RFC-2: PD 分离（Prefill-Decode Disaggregation） ​

1. 概述 ​

2. 背景与现状 gap ​

2.1 现状 vs sglang / tpu-inference 对比 ​

2.2 JAX backend 约束 ​

2.3 部署形态选择见 §4 ​

3. 架构总览 ​

4. 部署形态：SPMD vs MPMD ​

4.1 概念 ​

4.2 物理 / 软件层约束（为什么 MPMD 即便同 slice 也走 host） ​

4.3 Trade-off 与 sgl-jax 选择 ​

4.4 未来 SPMD 适用场景（§9 范围外） ​

4.5 Cross-process 部署细节（第一版实施） ​

5. 模块设计 ​

5.1 4-tuple ABC ​

5.2 JaxTransferKVEngine 数据链路 ​

5.2.1 P 端 KVSender.send（模式 A） ​

5.2.2 D 端 KVReceiver.pull（模式 A） ​

5.2.3 其他约定 ​

5.3 BootstrapServer（HTTP） ​

5.4 Scheduler Mixin ​

5.5 D2H staging（第一版 off） ​

5.6 接口预留：多 pool transfer descriptor ​

5.7 MP+ND ​

6. 关键数据流 ​

6.1 P → D KV 传输（模式 A，默认） ​

6.2 Proxy 注入字段 schema（顶层 + 子字段两层） ​

6.3 PD + ChunkCache（D 视角） ​

7. 配置项 ​

8. 实施路线 ​

9. 范围外 / 未来工作 ​

10. 未决问题 ​

RFC-2: PD 分离（Prefill-Decode Disaggregation）

1. 概述

2. 背景与现状 gap

2.1 现状 vs sglang / tpu-inference 对比

2.2 JAX backend 约束

2.3 部署形态选择见 §4

3. 架构总览

4. 部署形态：SPMD vs MPMD

4.1 概念

4.2 物理 / 软件层约束（为什么 MPMD 即便同 slice 也走 host）

4.3 Trade-off 与 sgl-jax 选择

4.4 未来 SPMD 适用场景（§9 范围外）

4.5 Cross-process 部署细节（第一版实施）

5. 模块设计

5.1 4-tuple ABC

5.2 JaxTransferKVEngine 数据链路

5.2.1 P 端 `KVSender.send`（模式 A）

5.2.2 D 端 `KVReceiver.pull`（模式 A）

5.2.3 其他约定

5.3 BootstrapServer（HTTP）

5.4 Scheduler Mixin

5.5 D2H staging（第一版 off）

5.6 接口预留：多 pool transfer descriptor

5.7 MP+ND

6. 关键数据流

6.1 P → D KV 传输（模式 A，默认）

6.2 Proxy 注入字段 schema（顶层 + 子字段两层）

6.3 PD + ChunkCache（D 视角）

7. 配置项

8. 实施路线

9. 范围外 / 未来工作

10. 未决问题