调研：sglang 与 tpu-inference 的 PD 分离

目的：为 sgl-jax 实现 PD（Prefill-Decode）分离提供上游对照。控制面以 sglang 为主参考（4-tuple ABC + Mixin 模式 + Bootstrap），数据面以 tpu-inference 的 JAX/Pallas API 为主参考（jax.experimental.transfer + Pallas DMA + ZMQ side channel）。
代码基准：
sglang：origin/main HEAD f04c52253（2026-05-18）
tpu-inference：main HEAD（2026-05-18）
与现有文档的关系：
gap-pd-disaggregation.md：tpu-inference 单边详解（含 In-process 三线程编排器、Pallas DMA 代码、设备分配机制等完整细节）
本文：sglang ↔ tpu-inference 双边对照（按"控制面 / 数据面 / 与 cache 关系"分章节，强调 sgl-jax 适配点）
配套：多级 KV cache 调研见 sglang-tpu-inference-multi-level-cache.md。

摘要 / TL;DR

议题	sglang 现状	tpu-inference 现状
部署形态	默认跨进程（每个 P/D 节点独立进程）	双形态：In-process（单进程多 mesh，开发调试用）+ Cross-process（`jax.experimental.transfer`，生产）
控制面 ABC	4-tuple：`BaseKVManager` / `BaseKVSender` / `BaseKVReceiver` / `BaseKVBootstrapServer`	自定义 `TPUConnectorScheduler` / `TPUConnectorWorker`
状态机	`KVPoll`（Failed / Bootstrapping / WaitingForInput / Transferring / Success）	隐含在 `_DisaggOrchestrator` 三线程流转中
Backend 数	5 个（Mooncake / NIXL / MoRI / Ascend / Fake）	1 个（`jax.experimental.transfer`）
Backend 注册	硬编码 enum + factory（无运行时注册）	—
Bootstrap Server	HTTP-based（`CommonKVBootstrapServer`）	HTTP Proxy + ZMQ side channel（无独立 bootstrap）
Scheduler 集成	Mixin 模式（`SchedulerDisaggregation{Prefill,Decode}Mixin`）	Engine 双实例 + 三线程编排器（仅 In-process）
跨进程数据路径	RDMA（GPUDirect 可绕 host）	必经 host 内存（无 GPUDirect 对应物）
D2H staging	`cudaMemcpyAsync` + pinned host	Pallas `copy_to_host` → 预分配 pinned host pool → `await_pull`（默认单机；`TPU_ENABLE_D2H_TRANSFER` 控制）
Partial KV pulling	RDMA 支持（按需拉部分）	不支持（`jax.experimental.transfer` 限制）
Tree cache 关系	PD 主路径绕过 tree_cache（直接调 `KVSender.send`）	—
MPMD vs SPMD	每 TP rank 独立进程；KV manager 每 rank 一个	SPMD 单 scheduler + thread-based worker

sgl-jax 实施方向（第一版）：直接 Cross-process（生产场景），不实现 In-process（详见 §一）；唯一 backend = jax.experimental.transfer；D2H staging：tpu-inference 上游单机默认开启、跨机强制关闭；sgl-jax 第一版统一关闭（最小工程量起点，详见 RFC-2 §5.5）；接口预留 spec decode 兼容性（详见 [[rfc_2_pd_disaggregation]]）。

零、术语与全景

0.1 术语

术语	含义
Prefill (P)	处理整个 prompt → 输出第一个 token + KV cache
Decode (D)	基于 KV cache 自回归生成后续 tokens
PD 分离	P 和 D 在物理/进程层面解耦，KV cache 显式传递
Bootstrap Server	协调 P/D 节点会合（注册、地址交换、room 号分配）
bootstrap_room	每个请求的唯一会合标识（P 注册，D 查询）
KV Connector	KV 传输的抽象接口（sglang 4-tuple）
Transfer Engine	底层传输引擎（Mooncake / NIXL / `jax.experimental.transfer` 等）
Proxy Server	HTTP 层请求路由（先 P → 提取 transfer_params → 转发 D）

0.2 全景模块图

text

  ┌────────────────────────────────────────────────────────────────────┐
  │                         HTTP 层（Proxy）                              │
  │  Client ──→ Proxy ──→ P 节点 (max_tokens=1, 仅 prefill)              │
  │                  │      ↓ 提取 kv_transfer_params                    │
  │                  └──→ D 节点 (stream=True, 附 transfer_params)        │
  └────────────────────────────────────────────────────────────────────┘
                                  │
  ┌────────────────────────────────▼─────────────────────────────────┐
  │                         控制面（参考 sglang）                       │
  │                                                                  │
  │  ┌───────────────────────┐         ┌───────────────────────┐    │
  │  │ Scheduler (P 节点)     │         │ Scheduler (D 节点)     │    │
  │  │  + PrefillMixin       │         │  + DecodeMixin        │    │
  │  │   - bootstrap_queue   │         │   - prealloc_queue    │    │
  │  │   - inflight_queue    │         │   - transfer_queue    │    │
  │  │   - event_loop_       │         │   - event_loop_       │    │
  │  │     {normal,overlap}_ │         │     {normal,overlap}_ │    │
  │  │     disagg_prefill    │         │     disagg_decode     │    │
  │  └──────────┬────────────┘         └──────────┬────────────┘    │
  │             │                                  │                 │
  │  ┌──────────▼──────────────┐       ┌──────────▼──────────────┐  │
  │  │ KVManager (P)            │       │ KVManager (D)            │  │
  │  │  KVSender                │       │  KVReceiver              │  │
  │  │  ┌────────────────────┐ │       │  ┌────────────────────┐ │  │
  │  │  │ init()              │ │       │  │ init()              │ │  │
  │  │  │ send(indices, ...)  │ │       │  │ send_metadata(...)  │ │  │
  │  │  │ poll() → KVPoll     │ │       │  │ poll() → KVPoll     │ │  │
  │  │  └────────────────────┘ │       │  └────────────────────┘ │  │
  │  └──────────┬──────────────┘       └──────────┬──────────────┘  │
  │             │                                  │                 │
  │             └──────────┬───────────────────────┘                 │
  │                        │                                          │
  │                ┌───────▼──────────┐                               │
  │                │ Bootstrap Server  │ (HTTP / ZMQ)                  │
  │                │  • 注册 P info     │                               │
  │                │  • D 查询 P 地址  │                                │
  │                │  • room 号管理    │                                │
  │                └──────────────────┘                                │
  └────────────────────────────────┬─────────────────────────────────┘
                                   │
  ┌─────────────────────────────────▼────────────────────────────────┐
  │              数据面（参考 tpu-inference / JAX API）                 │
  │                                                                  │
  │  Layer 1: jax.experimental.transfer（P2P，必经 host）              │
  │    P: kv_transfer_server.await_pull(uuid, kv)                     │
  │    D: conn.pull(uuid, kv_spec) → 落到本地 sharding                │
  │                                                                  │
  │  Layer 2: Pallas DMA（HBM ↔ HBM scatter）                         │
  │    D 端拉到 staging 后用 multi_layer_copy scatter 到 paged pool   │
  │                                                                  │
  │  Layer 3: ZMQ side channel（控制信息）                             │
  │    D 拉取完成 → 通知 P 释放 KV 缓冲区                                │
  └──────────────────────────────────────────────────────────────────┘

一、两种部署形态

sgl-jax 实施方向：第一版只实现 Cross-process（生产场景）。In-process 仅作为 tpu-inference 的已有形态做背景说明，sgl-jax 不实现——理由见 §1.3 末尾。

1.1 In-process（Local Disagg）：单进程多 mesh — tpu-inference 特有，sgl-jax 不实现

仅 tpu-inference 有此形态。sglang 默认是跨进程模型。

适用场景：单机多 chip 开发调试（如 TPU v5e 8 chip → 4 prefill + 4 decode），在同一 JAX runtime 内通过 jax.device_put / jax.experimental.reshard 完成 KV 跨 mesh 传输。基础设施由 DisaggExecutor（解析 PREFILL_SLICES=4 DECODE_SLICES=4 环境变量 + 分配双 mesh）+ _DisaggOrchestrator（三类线程 + 有界队列 + JetThread）支撑。

完整实现细节（设备分配、三线程编排、JetThread 崩溃处理、multi_layer_copy Pallas DMA 等）见 gap-pd-disaggregation.md §2.1 — 本文不再重复，因为这条路径 sgl-jax 不会落地。

关键约束：仅当 P 和 D 在同一 JAX runtime 内才可用。多进程 PD 下，每个进程的 jax.devices() 只包含本进程的芯片子集，无法做跨 runtime 的 ICI 传输——必须走 jax.experimental.transfer（即 §1.2 Cross-process）。

1.2 Cross-process（Multi-host Disagg）：跨进程 — sgl-jax 第一版方向

sglang 标准模式：

4-tuple ABC + 5 个具体 backend（Mooncake / NIXL / MoRI / Ascend / Fake）
HTTP Bootstrap Server 协调
Mixin 模式分隔 Scheduler 事件循环

tpu-inference 实现：

jax.experimental.transfer 作为唯一 backend
ZMQ 作为 pull-done 通知 side channel
HTTP Proxy 路由请求

sgl-jax 第一版的 PD 实现完全采用此形态。详细 API 见 §四；控制面接口见 §二/§三。

1.3 形态对比与 sgl-jax 取舍

维度	In-process	Cross-process
进程数	1	N (P) + M (D) + 1 (Proxy/Bootstrap)
KV 跨 mesh API	`jax.device_put` / `experimental.reshard`	`jax.experimental.transfer`
数据路径	进程内（HBM 间，ICI 直连）	必经 host 内存（DCN）
Bootstrap 协议	不需要（直接队列）	HTTP 注册 + 查询
故障传播	单进程，OOM/crash 影响整体	隔离好
部署复杂度	低	中（需要 Proxy + Bootstrap）
生产可用	✗	✓
sgl-jax 第一版	不实现	✅ 直接做

为什么 sgl-jax 跳过 In-process：

In-process 的复杂度集中在基础设施（DisaggExecutor 多 mesh 分配 / _DisaggOrchestrator 三线程 + 有界队列 / JetThread 崩溃即 SIGKILL），这些几乎全部不能复用到 Cross-process——Cross-process 用进程隔离 + Bootstrap + transfer server 替代
真正能跨形态复用的只有"按 block_id 提取 KV"和"scatter KV 到指定 block"两个小模块，且这两个在 Cross-process 框架下用单进程跑单元测试就能验证
In-process 不属于生产场景，投入工作量换调试便利不划算
直接做 Cross-process 单机部署（同机器跑 P 进程 + D 进程 + Proxy + Bootstrap）已经能覆盖大部分调试需求

二、控制面：sglang 4-tuple ABC + Bootstrap

2.1 4-tuple Base Abstraction

python/sglang/srt/disaggregation/base/conn.py：

python

class KVArgs:                                  # L30, POD
    kv_data_ptrs: List[int]                   # GPU pool 指针
    kv_data_lens: List[int]
    item_lens: List[int]
    aux_data_ptrs / aux_data_lens / aux_item_lens
    state_data_ptrs / state_data_lens / state_type: "none"|"mamba"|"swa"
    gpu_id: int                                # ← TPU 迁移要替换（用 device_id 或 sharding 描述）
    pp_rank: int
    prefill_start_layer: int
    system_dp_rank: int

class KVPoll:                                  # L71, 状态机
    Failed = 0
    Bootstrapping = 1
    WaitingForInput = 2
    Transferring = 3
    Success = 4

class BaseKVManager(ABC):                      # L79
    def __init__(self, args: KVArgs, mode: DisaggregationMode,
                 server_args: ServerArgs, is_mla_backend: bool): ...
    def register_to_bootstrap(self): ...

class BaseKVSender(ABC):                       # L97
    def __init__(self, mgr, bootstrap_addr, bootstrap_room,
                 dest_tp_ranks, pp_rank): ...
    def init(self, num_kv_indices, aux_index): ...
    def send(self, kv_indices, state_indices): ...
    def poll(self) -> KVPoll: ...
    def failure_exception(self) -> Exception: ...

class BaseKVReceiver(ABC):                     # L153
    def __init__(self, mgr, bootstrap_addr, bootstrap_room): ...
    def init(self, prefill_dp_rank): ...
    def send_metadata(self, kv_indices, aux_index, state_indices): ...
    def poll(self) -> KVPoll: ...
    def clear(self): ...
    def abort(self): ...

class BaseKVBootstrapServer(ABC):              # L213
    def __init__(self, host, port): ...

（注：辅助类 StateType L17、KVTransferMetric L24 在同文件中。）

2.2 Backend 继承层级

text

BaseKVManager / BaseKVSender / BaseKVReceiver / BaseKVBootstrapServer  ← base/conn.py
    │
    └── CommonKV*  ← common/conn.py (Manager L90 / Sender L606 / Receiver L731 / BootstrapServer L934)
            ├── MooncakeKV*  ← mooncake/conn.py (Manager L184 / Sender L1657 / Receiver L1758 / BootstrapServer L1921)
            │       └── AscendKV*  ← ascend/conn.py
            ├── NixlKV*      ← nixl/conn.py
            └── MoRIKV*      ← mori/conn.py

FakeKV*  ← fake/conn.py（直接继承 Base，无 bootstrap）

CommonKVManager 是真正的"基础设施"层（common/conn.py:90）：集中了 bootstrap 注册、TP/PP rank mapping、KV ptr 计算（get_mha_kv_ptrs_with_pp L441、get_mla_kv_ptrs_with_pp L474）。新 backend 通常继承 Common*。

2.3 Factory & Role Assignment

python/sglang/srt/disaggregation/utils.py：

python

class DisaggregationMode(Enum):                # L33
    NULL = "null"
    PREFILL = "prefill"
    DECODE = "decode"

class TransferBackend(Enum):                   # L312
    MOONCAKE / NIXL / MORI / ASCEND / FAKE

class KVClassType(Enum):                       # L320
    MANAGER / SENDER / RECEIVER / BOOTSTRAP_SERVER

def get_kv_class(backend: TransferBackend, class_type: KVClassType):  # L329（含多个 overload L333/L337/L341/L345）
    # 5-way dispatch (mooncake/mori/ascend/nixl/fake)

Role 判定：进程级，由 server_args.disaggregation_mode 字符串决定（与 sglang scheduler 中的事件循环 dispatch 一一对应）。

2.4 Bootstrap Server 协议

HTTP-based，入口 common/conn.py:934：

python

class CommonKVBootstrapServer:                     # L934
    def run(self): ...                             # L962
    def _setup_routes(self): ...                   # L979

# Prefill 侧调用
CommonKVManager.register_to_bootstrap()            # L376
# 注册 PrefillServerInfo (L50) / PrefillRankInfo (L81)

# Decode 侧调用
CommonKVReceiver._get_bootstrap_info_from_server() # L844
CommonKVReceiver.query_prefill_dp_ranks()          # L864

bootstrap_room 由 receiver 在 __init__ 时 caller-supplied（BaseKVReceiver.__init__ 第三个参数），通常是 64-bit 整数。

重要观察：没有运行时注册机制——新 backend 必须改 utils.py 加 enum 分支。sgl-jax 可以做更优雅的注册表（例如 entry_points 或 decorator 注册）。

2.5 警告：`sglang.connector.BaseKVConnector` ≠ PD ABC

易混淆点：sglang 的 python/sglang/srt/connector/base_connector.py：

python

class BaseConnector(ABC):     # L13
    def weight_iterator(self): ...
    def pull_files(self): ...

class BaseKVConnector(BaseConnector):    # L75
    def get / set / getstr / setstr / list

这是 model weight loading 用的（Redis 等存储后端），不是 vLLM 风格的 PD 插件抽象。

sgl-jax 的含义：

设计 PD 时不要照搬 vLLM v1 的 KVConnectorBase_V1 名字（容易和 sglang 的 BaseKVConnector 冲突）
可以借鉴 vLLM v1 的「scheduler-worker 双 connector」风格，但接口签名应对齐 sglang 的 4-tuple

三、控制面：Scheduler 集成（Mixin 模式）

3.1 Prefill 节点

Scheduler 字段（scheduler.py 中 P 节点初始化区域）：

python

self.disagg_prefill_bootstrap_queue = PrefillBootstrapQueue(...)
self.disagg_prefill_inflight_queue: List[Req] = []

Mixin 方法（disaggregation/prefill.py:358 class SchedulerDisaggregationPrefillMixin）：

方法	行号	用途
`event_loop_normal_disagg_prefill`	392	正常模式事件循环
`event_loop_overlap_disagg_prefill`	425	overlap 模式事件循环
`process_batch_result_disagg_prefill`	469	处理 batch 完成 → 加入 inflight queue
`process_disagg_prefill_inflight_queue`	597	轮询 sender.poll()，状态机推进
`process_prefill_chunk`	727	处理 chunked prefill
`send_kv_chunk`	755	实际发送 KV（绕过 tree_cache，直读 req_to_token_pool）

PrefillBootstrapQueue 定义在 disaggregation/prefill.py:91。

3.2 Decode 节点

Scheduler 字段（scheduler.py 中 D 节点初始化区域）：

python

self.disagg_decode_transfer_queue = DecodeTransferQueue(...)
self.disagg_decode_prealloc_queue = DecodePreallocQueue(...)

队列实现：

类	文件	行号
`PrefillBootstrapQueue`	`disaggregation/prefill.py`	91
`DecodePreallocQueue`	`disaggregation/decode.py`	260
`DecodeTransferQueue`	`disaggregation/decode.py`	1345
`DecodeReqToTokenPool`	`disaggregation/decode.py`	102（PD 专用）
`HybridMambaDecodeReqToTokenPool`	`disaggregation/decode.py`	179
`DecodeRequest`	`disaggregation/decode.py`	245

Mixin 方法（disaggregation/decode.py:1572 class SchedulerDisaggregationDecodeMixin）：

方法	行号	用途
`event_loop_normal_disagg_decode`	1574	正常模式事件循环
`event_loop_overlap_disagg_decode`	1601	overlap 模式事件循环
`get_next_disagg_decode_batch_to_run`	1649	选 batch
`get_new_prebuilt_batch`	1683	取已传输完成的请求
`process_decode_queue`	1758	轮询 receiver.poll()
`_commit_transfer_to_req`	1383（属于 `DecodeTransferQueue`）	写 metadata_buffers（不调 tree_cache.insert）

3.3 事件循环 dispatch 矩阵

python

if disaggregation_mode == NULL:
    if enable_pdmux:        event_loop_pdmux()
    elif pp_size > 1:       event_loop_pp()
    elif enable_overlap:    event_loop_overlap()
    else:                   event_loop_normal()
elif disaggregation_mode == PREFILL:
    if pp_size > 1:         event_loop_pp_disagg_prefill()
    elif enable_overlap:    event_loop_overlap_disagg_prefill()
    else:                   event_loop_normal_disagg_prefill()
elif disaggregation_mode == DECODE:
    if pp_size > 1:         event_loop_pp_disagg_decode()
    elif enable_overlap:    event_loop_overlap_disagg_decode()
    else:                   event_loop_normal_disagg_decode()

意义：Mixin 模式让 PD 逻辑不污染 Scheduler 主类——sgl-jax 强烈建议沿用。

3.4 GPU MPMD vs TPU SPMD 对 Mixin 的影响

维度	sglang / GPU（MPMD）	sgl-jax / TPU（SPMD）
Scheduler 进程数	每 TP rank 一个	单进程
`KVManager` 实例数	每 rank 一个（各自 register_to_bootstrap）	单实例（per-process singleton）
Bootstrap 注册粒度	per-rank（每 rank 上报自己的地址）	per-process（单地址）
Worker 通信	进程间 IPC	线程间共享内存
失败传播	一 rank 失败可能不阻塞其他 rank	单进程失败即整体失败
Mixin 实装难度	中（需处理 rank 同步）	低（直接照搬，事件循环单线程）

关键事实：sgl-jax 的 Mixin 实装比 sglang 更简单，因为不需要处理 rank 间同步。

四、数据面：tpu-inference JAX/Pallas API

4.1 跨进程传输：`jax.experimental.transfer`

4.1.1 实验性 API 警告

jax.experimental.transfer 是内部实验性 API，截至 2026-05 几乎没有公开文档。所有用法基于 tpu-inference 仓库的代码行为和 TPU 硬件架构推断。

Risk：

API 签名可能在 JAX 版本升级时改变
错误处理路径未文档化
多 host 并发拉取的行为细节未文档化

sgl-jax 应对：

封装一层自己的 ABC（如 KVTransferEngine），底层默认走 jax.experimental.transfer
短期内只支持单机 PD（避免触发未知多 host 路径）
准备 fallback：用 ZMQ + 显式 jax.device_get / device_put 实现一个备用 backend

4.1.2 Server 启动（Decode 端）

python

# distributed/tpu_connector.py:573
from jax.experimental.transfer import start_transfer_server

self.kv_transfer_server = start_transfer_server(
    jax.local_devices()[0].client,
    server_addr,                                       # 控制地址 "ip:port"
    [transport_addr] * get_transfer_channel_number(), # 数据通道列表（默认 8 个）
    max_num_parallel_copies=8,
    transfer_size=256 * 1024 * 1024,                   # 256 MB
    use_raw_buffers=False,
)

参数说明：

参数	说明
`client`	从 `jax.local_devices()[0].client` 拿到的 PJRT 客户端
`server_addr`	TCP `"host_ip:port"` 控制地址
`transport_addr` 列表	数据通道地址列表（每个一条 TCP 连接）；由 `TPU_KV_TRANSFER_CHANNEL_NUMBER` 控制
`transfer_size`	单次传输 chunk 大小，256 MB
`max_num_parallel_copies`	并发拷贝数，8
`use_raw_buffers`	是否绕过 JAX 的高层抽象（实验性）

4.1.3 Producer 侧（P 节点）

python

# tpu_connector.py:710 / 760
self.kv_transfer_server.await_pull(uuid, kv_data)

参数：

uuid：唯一传输 ID（用 Python uuid.uuid4() 生成）
kv_data：要发布的 jax.Array（或 list[Array]）；可以在 HBM 上也可以在 pinned host 上
阻塞直到 Consumer pull 完成或超时

4.1.4 Consumer 侧（D 节点）

python

# tpu_connector.py:772 - 建立连接
conn = self.kv_transfer_server.connect(f"{remote_host}:{remote_port}")

# tpu_connector.py:793 - 拉取数据
kv_spec = [
    jax.ShapeDtypeStruct(
        shape=(num_blocks, *cache_inner_shape),
        dtype=dtype,
        sharding=self.sharding,   # ← device sharding → 数据直接落到 HBM
    )
    for _ in range(num_layers)
]
kv = conn.pull(uuid, kv_spec)

关键：kv_spec.sharding 决定数据落在哪儿：

device sharding：数据直接到 Consumer 的 HBM
host sharding：落到 pinned host memory

坑点：

不支持 partial pulling（不像 RDMA）：即使 D 端有 prefix cache 命中，也必须拉取全部 prefill block
不支持 native async pulling：当前用 ThreadPoolExecutor + is_ready() 轮询模拟
拉取完成检查：chunk.is_ready()（间隔 1ms 轮询），非 block_until_ready

4.2 数据路径必经 host：无 GPUDirect RDMA 对应物

架构 2 PD 分离不能做真正的 device-to-device 直传——即使同一 pod 内芯片有 ICI 物理直连，由于 P 和 D 处于不同 JAX runtime，必须经过 host 内存：

text

P TPU Chip (HBM)
       │ PCIe
       ↓
P Host Memory (pinned_host)
       │ DCN / TCP
       ↓
D Host Memory
       │ PCIe
       ↓
D TPU Chip (HBM)

与 GPU 对比：GPU 有 GPUDirect RDMA（NIC 直接读写 GPU VRAM 绕过 host），TPU 没有对应物。这是硬件约束，不是软件设计选择。

RFC 中应明确：

PD 分离的网络延迟比 GPU 高一截
多 host 部署时这是主要瓶颈
缓解：增大 batch、减少传输次数、用更高带宽 DCN

4.3 D2H staging：`await_pull` 的可选优化层（不是替代）

澄清易混点：D2H staging 与 jax.experimental.transfer 不是两种 D2H 实现方式的选择——前者是后者之前的一步可选优化。

4.3.1 完整数据路径

text

                  P 端 HBM 上的 KV
                       │
        ┌──────────────┴───────────────┐
        │ (路径 A，默认单机)             │ (路径 B，多 host 或 D2H staging=off)
        ▼                              ▼
  Pallas copy_to_host             [跳过 staging]
        │                              │
        ▼                              │
  Pinned host pool (HostKVPool, 预分配) │
        │                              │
        └──────────────┬───────────────┘
                       │
                       ▼
        await_pull(buffer)   ← jax.experimental.transfer
                       │
                       ▼
                   网络 (DCN)
                       │
                       ▼
                D 端 conn.pull(uuid, kv_spec)
                       │  (kv_spec.sharding 永远是 device sharding)
                       ▼
                D 端 HBM 上的 KV (库内部隐含 H2D，对调用方透明)

4.3.2 关键事实

事实	来源
`jax.experimental.transfer.await_pull` 接受任意 sharding 的 `jax.Array` ——理论上可以直接传 HBM array	`jax/experimental/transfer.py` 签名
tpu-inference 在 `TPU_ENABLE_D2H_TRANSFER=true`（默认）时做显式 staging：HBM → Pallas `copy_to_host` → 预分配的 pinned host pool（`HostKVPool`） → `await_pull(host_buffer)`	`tpu_connector.py:540-545, 695-745`, `host_kv_pool.py:53-60`
Staging 的目的：消除运行时 PJRT `MapDmaBuffer` 的系统调用开销 + 把 HBM 占用尽早释放降低 device 内存压力	`tpu_connector.py:715` 注释 "reduce the memory pressure on device"
多 host 场景下 staging 路径未实现（路径 A 不适用），由 transfer 库内部处理 D2H	tpu-inference 代码行为推断（未文档化）
D 端从不做 H2D staging：`_get_kv_spec` 永远返回 device sharding，`conn.pull` 直接产出 HBM Array；H2D 由 transfer 库内部完成，对调用方透明	`tpu_connector.py:788-793, 830-836`

4.3.3 `TPU_ENABLE_D2H_TRANSFER` flag 控制的两种模式

模式	P 端行为	目的
true（默认，单机）	显式 Pallas `copy_to_host` → pinned host pool → `await_pull(host_buffer)`	用 `HostKVPool` 预分配 pinned host buffer，消除运行时 `MapDmaBuffer` 开销；缓解 device 内存压力
false 或多 host	直接 `await_pull(hbm_kv)`	多 host 场景下 staging 路径未实现；或希望由 transfer 库内部自己处理 D2H

4.3.4 Pallas `copy_to_host` 与 HiCache 的 `jax.device_put` 的不同定位

路径	API	何时用
HiCache offload（详见 cache 文档 §2.2）	`jax.device_put(arr, host_sharding)` + `block_until_ready`	后台异步、不在主线程关键路径；实现简单
PD D2H staging	Pallas `copy_to_host(src, dest_host_buffer)`	请求关键路径、延迟敏感；目标是预分配的 pinned host buffer 而非每次新分配

为什么 PD 用 Pallas 而非 jax.device_put：

jax.device_put 每次创建新的 jax.Array（PJRT 内部 pool 走分配/释放），有 MapDmaBuffer 系统调用开销
Pallas copy_to_host 直接操作 DMA 引擎，写入已预分配的 HostKVPool buffer，零分配开销
Pallas 配合 is_ready() 轮询用作传输完成检查（更细粒度），比 block_until_ready 适合关键路径

Pallas copy_to_host 实现：

python

# distributed/kv_transfer.py:416
def copy_to_host(src: jax.Array, dest: jax.Array, ...) -> jax.Array:
    # Pallas kernel 内部：
    # async_copy = pltpu.make_async_copy(src_ref, dest_ref_in_pltpu.HOST, semaphore_ref)
    # async_copy.start(); async_copy.wait()

注：这里 dest 必须是 HostKVPool 中预分配的 pinned host slot，由 await_pull 在 staging 完成后接管发送。

4.4 传输参数调优

环境变量	默认	说明
`TPU_KV_TRANSFER_PORT`	9100	P2P 控制端口
`TPU_SIDE_CHANNEL_PORT`	9600	ZMQ 通知端口
`TPU_KV_TRANSFER_CHANNEL_NUMBER`	8	并行传输通道数
`TPU_ENABLE_D2H_TRANSFER`	true	启用 P 端显式 D2H staging（详见 §4.3.3）
`TPU_P2P_WAIT_PULL_TIMEOUT`	30	缓冲区保持秒数
`transfer_size`（`start_transfer_server` 参数）	256MB	单次传输 chunk 大小
`max_num_parallel_copies`（同上）	8	并发拷贝数

4.5 ZMQ side channel：pull-done 通知

用途：D 端拉取完成后通过 ZMQ 通知 P 端释放为该传输保留的 KV 缓冲区。

python

# tpu_connector.py: P 端启动 ZMQ ROUTER socket 监听
def _start_zmq_listener(self):
    zmq_port = int(os.environ.get("TPU_SIDE_CHANNEL_PORT", 9600))
    context = zmq.Context()
    self.zmq_socket = context.socket(zmq.ROUTER)
    self.zmq_socket.bind(f"tcp://*:{zmq_port}")
    # 后台线程监听
    self.zmq_thread = threading.Thread(
        target=self._pull_notify_listener, daemon=True
    )
    self.zmq_thread.start()

def _pull_notify_listener(self):
    while True:
        identity, _, message = self.zmq_socket.recv_multipart()
        notify = json.loads(message)
        transfer_uuid = notify["transfer_uuid"]
        # 释放 P 端为该传输保留的 KV 缓冲区
        self._release_kv_buffer(transfer_uuid)

# D 端：通过 ZMQ DEALER 通知 P 拉取完成
def _notify_pull_complete(self, transfer_uuid: str):
    context = zmq.Context()
    socket = context.socket(zmq.DEALER)
    socket.connect(f"tcp://{self.prefill_addr}:{self.zmq_port}")
    socket.send_json({"transfer_uuid": transfer_uuid})
    socket.close()

与 sglang Bootstrap 的关系：

sglang 的 Bootstrap Server 处理会合（P 注册地址、D 查询地址、room 号分配）
ZMQ side channel 处理传输生命周期信号（pull 完成 / 释放）
两者职责不同，sgl-jax 中应分开实现

五、Pallas DMA：`multi_layer_copy`（HBM ↔ HBM scatter）

D 端从 P 拉取到 staging（pinned_host 或 device sharding）后，需要 scatter 到本地 paged KV pool 的指定 block 位置。这一步用 Pallas DMA 实现，跨 layer 流水线化。

5.1 跨 layer 流水线时序

text

Layer 0: start_copy ──────────────────────────────────
Layer 1:              start_copy ──────────────────────
Layer 0:                          wait_copy ───────────
Layer 2:                                    start_copy
Layer 1:                                    wait_copy
...

5.2 `pltpu.make_async_copy` + `SemaphoreType.DMA`

python

def multi_layer_copy(
    *,
    src_array: list[jax.Array],
    dest_array: list[jax.Array],
    src_offsets: jax.Array,
    dest_offsets: jax.Array,
    chunk_sizes: jax.Array,
    num_chunks: jax.Array | None = None,
    mesh: Mesh | None = None,
    src_sharding_spec: P | None = None,
    dest_sharding_spec: P | None = None,
    replicated_sharding_spec: P | None = None,
)

实现层级（distributed/kv_transfer.py）：

text

multi_layer_copy                  (L296)
  └─ _async_copy_jit              (L187, JIT 入口)
       └─ jax.shard_map 内部:
            ├─ _start_chunked_copy_kernel      ← Pallas kernel
            │    └─ pltpu.make_async_copy(...).start()
            │
            └─ _wait_for_chunked_copy_kernel   ← Pallas kernel
                 └─ pltpu.make_async_copy(...).wait()

5.3 `multi_layer_copy` vs `dynamic_update_slice_in_dim`

API	优势	劣势	何时用
`multi_layer_copy`（Pallas）	跨 layer 重叠、最大 DMA 带宽	实现复杂、Pallas 学习曲线	多 block + 跨 layer（HiCache load / PD scatter）
`dynamic_update_slice_in_dim`	简单纯 JAX	只更新连续 slice、无 layer 间流水	retract / single-block 操作

5.4 与 CUDA 自定义 scatter kernel 对应

SGLang (CUDA)	tpu-inference (Pallas)
自定义 CUDA scatter kernel（`sgl-kernel/csrc/kvcacheio/transfer.cu`）	`multi_layer_copy` (Pallas + `pltpu.make_async_copy`)
`cudaMemcpyAsync` (cross-stream)	多个 `make_async_copy` 自动重叠（XLA 调度）
CUDA semaphore	`pltpu.SemaphoreType.DMA`

六、Proxy Server 与请求路由

6.1 sglang 风格（`sglang_router/mini_lb.py`）

sglang 在 sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py 提供官方 Mini Load Balancer。启动命令：

bash

python -m sglang_router.launch_router \
  --pd-disaggregation \
  --prefill http://127.0.0.1:30000 \
  --decode http://127.0.0.1:30001 \
  --host 0.0.0.0 --port 8000

关键流程（mini_lb.py:358-405 + :135-141）：

lb.select_pair() 选 (P, D) 实例 + P 的 bootstrap_port

Proxy 在 req 顶层注入三字段（不是嵌套在 kv_transfer_params）：

python

modified_request.update({
    "bootstrap_host": hostname,                          # P 的 host
    "bootstrap_port": bootstrap_port,                    # P 的 bootstrap port
    "bootstrap_room": _generate_bootstrap_room(),        # Proxy 生成 (64-bit)
})

asyncio.gather 同时 post 给 P 与 D（不是串行）：

python

tasks = [
    session.post(f"{prefill_server}/{endpoint}", json=prefill_req),
    session.post(f"{decode_server}/{endpoint}", json=decode_req),
]
prefill_response, decode_response = await asyncio.gather(*tasks)

P/D 收到 req 后从顶层 bootstrap_host / bootstrap_port / bootstrap_room 字段读，由 tokenizer_manager.py:999-1001 透传到 scheduler。
dp_size > 1 时还会注入 routed_dp_rank + disagg_prefill_dp_rank（mini_lb 端 _fork_dp_requests）。

sgl-jax 实施方向：第一版 Proxy 直接复用 sglang_router/mini_lb（启动命令同上），不写自有 Proxy 代码——这是 sglang 主流做法，与 RFC-2 Proxy 部分对齐。

6.2 tpu-inference `examples/disagg/toy_proxy_server.py`

python

@app.post("/v1/completions")
async def completions(request: CompletionRequest):
    # 1. 发送到 Prefill server（max_tokens=1，仅做 prefill）
    prefill_response = await forward_to_prefill(
        request, max_tokens=1
    )

    # 2. 提取 KV transfer params
    kv_transfer_params = prefill_response["kv_transfer_params"]

    # 3. 转发到 Decode server（附带 transfer params，流式返回）
    decode_request = request.copy()
    decode_request.kv_transfer_params = kv_transfer_params

    return StreamingResponse(
        forward_to_decode_stream(decode_request)
    )

关键流程：

Client → Proxy 一个请求
Proxy 复制请求 → P 节点（强制 max_tokens=1，只做 prefill）
P 完成 → 返回 kv_transfer_params（含 transfer_uuid / prefill_addr / 等）
Proxy 把 kv_transfer_params 加入 Decode 请求 → D 节点
D 拉取 KV → 开始 decode → 流式返回给 Client

七、PD ↔ Tree Cache 的关系

7.1 PD 主路径绕过 tree_cache

sglang 的 PD 不通过 tree_cache 做 KV 传输：

python

# prefill.py:755+ (send_kv_chunk)
kv_indices = (
    self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx]
    .cpu()
    .numpy()
)
# ... 选 KV blocks 后
req.disagg_kv_sender.send(page_indices, state_indices)

直接读 req_to_token_pool、调 KVSender.send，与 tree 完全解耦。

D 端接收 KV 后不调 tree_cache.insert：

DecodeTransferQueue._commit_transfer_to_req (decode.py:1383) 只写 metadata_buffers
tree 插入由 D 完成 decode 后通过 cache_finished_req 走标准路径

7.2 ChunkCache + PD 一等公民

ChunkCache 是 PD 分离场景的合理选择——无 prefix 复用，纯 KV pool 管理，与 PD 传输模型对齐
sglang 默认 D 节点走 ChunkCache（实际上即使开了 enable_hierarchical_cache，PD 模式下 D 端通常也走 ChunkCache）
sgl-jax 第一版可只支持 ChunkCache + PD 组合（RFC-2 ADR-9）

7.3 UnifiedRadixCache + PD 可叠加

UnifiedRadixCache 完全不感知 PD（1960 行 0 处提及）
接口层正交：UnifiedRadixCache 与 PD 可以同时启用，但需要谨慎处理：
- P 端：完成 prefill → 既写 tree（用于本节点 prefix 复用）又通过 KVSender 发送给 D
- D 端：接收到完整 KV → 不立即写 tree（D 端 prefix 复用受 partial pull 限制收益低）

详细的耦合分析见 sglang-tpu-inference-multi-level-cache.md §5（30% 共享 / 70% 独立）。

参考文件索引

sglang（origin/main HEAD `f04c52253`）

文件	内容
`disaggregation/base/conn.py`	4-tuple ABC + `KVPoll` + `KVArgs`
`disaggregation/common/conn.py`	`CommonKV*`（基础设施层，bootstrap / rank mapping / KV ptr 计算）
`disaggregation/mooncake/conn.py`	Mooncake backend
`disaggregation/nixl/conn.py`	NIXL backend
`disaggregation/mori/conn.py`	MoRI backend
`disaggregation/ascend/conn.py`	Ascend backend（继承 Mooncake）
`disaggregation/fake/conn.py`	Fake backend（用于测试）
`disaggregation/prefill.py`	`PrefillBootstrapQueue` + `SchedulerDisaggregationPrefillMixin`
`disaggregation/decode.py`	`DecodePreallocQueue` / `DecodeTransferQueue` / `DecodeReqToTokenPool` / `SchedulerDisaggregationDecodeMixin`
`disaggregation/utils.py`	`DisaggregationMode` / `TransferBackend` / `KVClassType` enum + `get_kv_class` factory
`disaggregation/decode_kvcache_offload_manager.py`	唯一 PD↔HiCache 融合点

tpu-inference（main HEAD 2026-05-18）

文件	内容
`core/core_tpu.py`	`DisaggEngineCore`, `_DisaggOrchestrator`（三线程编排）
`core/disagg_executor.py`	`DisaggExecutor`（设备分配）
`core/disagg_utils.py`	`_parse_slices` 等工具函数
`distributed/tpu_connector.py`	`TPUConnectorScheduler` / `TPUConnectorWorker`（Multi-host P2P）
`distributed/kv_transfer.py`	`multi_layer_copy`（Pallas DMA）+ `copy_to_host`
`examples/disagg/toy_proxy_server.py`	HTTP Proxy 路由

配套文档

gap-pd-disaggregation.md — tpu-inference 单边 PD 详解（含 In-process 三线程编排器、Pallas DMA 代码完整版）
sglang-tpu-inference-multi-level-cache.md — 多级 KV cache 双边对照
sglang Issue #20415 — Unified Hybrid Radix Cache Refactor Roadmap

tpu-inference Q2 对照

调研：sglang 与 tpu-inference 的 PD 分离 ​

摘要 / TL;DR ​

零、术语与全景 ​

0.1 术语 ​

0.2 全景模块图 ​

一、两种部署形态 ​

1.1 In-process（Local Disagg）：单进程多 mesh — tpu-inference 特有，sgl-jax 不实现 ​

1.2 Cross-process（Multi-host Disagg）：跨进程 — sgl-jax 第一版方向 ​

1.3 形态对比与 sgl-jax 取舍 ​

二、控制面：sglang 4-tuple ABC + Bootstrap ​

2.1 4-tuple Base Abstraction ​

2.2 Backend 继承层级 ​

2.3 Factory & Role Assignment ​

2.4 Bootstrap Server 协议 ​

2.5 警告：sglang.connector.BaseKVConnector ≠ PD ABC ​

三、控制面：Scheduler 集成（Mixin 模式） ​

3.1 Prefill 节点 ​

3.2 Decode 节点 ​

3.3 事件循环 dispatch 矩阵 ​

3.4 GPU MPMD vs TPU SPMD 对 Mixin 的影响 ​

四、数据面：tpu-inference JAX/Pallas API ​

4.1 跨进程传输：jax.experimental.transfer ​

4.1.1 实验性 API 警告 ​

4.1.2 Server 启动（Decode 端） ​

4.1.3 Producer 侧（P 节点） ​

4.1.4 Consumer 侧（D 节点） ​

4.2 数据路径必经 host：无 GPUDirect RDMA 对应物 ​

4.3 D2H staging：await_pull 的可选优化层（不是替代） ​

4.3.1 完整数据路径 ​

4.3.2 关键事实 ​

4.3.3 TPU_ENABLE_D2H_TRANSFER flag 控制的两种模式 ​

4.3.4 Pallas copy_to_host 与 HiCache 的 jax.device_put 的不同定位 ​

4.4 传输参数调优 ​

4.5 ZMQ side channel：pull-done 通知 ​

五、Pallas DMA：multi_layer_copy（HBM ↔ HBM scatter） ​

5.1 跨 layer 流水线时序 ​

5.2 pltpu.make_async_copy + SemaphoreType.DMA ​

5.3 multi_layer_copy vs dynamic_update_slice_in_dim ​

5.4 与 CUDA 自定义 scatter kernel 对应 ​

六、Proxy Server 与请求路由 ​

6.1 sglang 风格（sglang_router/mini_lb.py） ​

6.2 tpu-inference examples/disagg/toy_proxy_server.py ​

七、PD ↔ Tree Cache 的关系 ​

7.1 PD 主路径绕过 tree_cache ​

7.2 ChunkCache + PD 一等公民 ​

7.3 UnifiedRadixCache + PD 可叠加 ​

参考文件索引 ​

sglang（origin/main HEAD f04c52253） ​

tpu-inference（main HEAD 2026-05-18） ​

配套文档 ​

调研：sglang 与 tpu-inference 的 PD 分离

摘要 / TL;DR

零、术语与全景

0.1 术语

0.2 全景模块图

一、两种部署形态

1.1 In-process（Local Disagg）：单进程多 mesh — tpu-inference 特有，sgl-jax 不实现

1.2 Cross-process（Multi-host Disagg）：跨进程 — sgl-jax 第一版方向

1.3 形态对比与 sgl-jax 取舍

二、控制面：sglang 4-tuple ABC + Bootstrap

2.1 4-tuple Base Abstraction

2.2 Backend 继承层级

2.3 Factory & Role Assignment

2.4 Bootstrap Server 协议

2.5 警告：`sglang.connector.BaseKVConnector` ≠ PD ABC

三、控制面：Scheduler 集成（Mixin 模式）

3.1 Prefill 节点

3.2 Decode 节点

3.3 事件循环 dispatch 矩阵

3.4 GPU MPMD vs TPU SPMD 对 Mixin 的影响

四、数据面：tpu-inference JAX/Pallas API

4.1 跨进程传输：`jax.experimental.transfer`

4.1.1 实验性 API 警告

4.1.2 Server 启动（Decode 端）

4.1.3 Producer 侧（P 节点）

4.1.4 Consumer 侧（D 节点）

4.2 数据路径必经 host：无 GPUDirect RDMA 对应物

4.3 D2H staging：`await_pull` 的可选优化层（不是替代）

4.3.1 完整数据路径

4.3.2 关键事实

4.3.3 `TPU_ENABLE_D2H_TRANSFER` flag 控制的两种模式

4.3.4 Pallas `copy_to_host` 与 HiCache 的 `jax.device_put` 的不同定位

4.4 传输参数调优

4.5 ZMQ side channel：pull-done 通知

五、Pallas DMA：`multi_layer_copy`（HBM ↔ HBM scatter）

5.1 跨 layer 流水线时序

5.2 `pltpu.make_async_copy` + `SemaphoreType.DMA`

5.3 `multi_layer_copy` vs `dynamic_update_slice_in_dim`

5.4 与 CUDA 自定义 scatter kernel 对应

六、Proxy Server 与请求路由

6.1 sglang 风格（`sglang_router/mini_lb.py`）

6.2 tpu-inference `examples/disagg/toy_proxy_server.py`

七、PD ↔ Tree Cache 的关系

7.1 PD 主路径绕过 tree_cache

7.2 ChunkCache + PD 一等公民

7.3 UnifiedRadixCache + PD 可叠加

参考文件索引

sglang（origin/main HEAD `f04c52253`）

tpu-inference（main HEAD 2026-05-18）

配套文档