docs(design): D→P RDMA snapshot push design

Goal: skip P-side re-prefill on reseed path. Push session KV
snapshot from D back to P after each direct-to-D append; reseed
re-uses P's snapshot to fire only the P→D' transfer (no model.forward
on P).

Decision: Option C — D→P snapshot at append-commit, P-side
PrefillSnapshotStore (side-table, not in radix tree), prefill
bypass when snapshot is fresh. Rejects A (radix multi-producer),
B (D→D' direct, fails for session-not-resident), D (eviction-only).

Lays out 8-commit roadmap, wire protocol, failure modes, and the
E4 experiment plan (KVC + D→P vs naive PD-disagg E1 baseline).
This commit is contained in:
Claude Code Agent
2026-05-13 00:44:03 +08:00
parent 6d1c9237fa
commit 9c35eddc79

View File

@@ -0,0 +1,446 @@
# D→P KV 反向推送设计
**日期**2026-05-12
**分支**`h200-cu130`(在此分支上做,后续 cherry-pick 到 `feat/d-to-p-sync` 备用)
**目标**:让 reseed 路径绕过 P 端 re-prefill把 reseed 总耗时从 3-7s 压到接近一次 RDMA P→D' 传输(~200-400ms
**前置**`docs/RESEED_SLOW_PATH_AND_D_TO_P_GAP_ZH.md`reseed 现状),`docs/KVC_EVICTION_GRANULARITY_DESIGN_ZH.md`(架构层背景)
---
## 0. TL;DR
1. **现状**v2 reseed 路径 = P open session + P 完整 re-prefill~1.5-3s+ P→D' mooncake transfer~200-400ms RDMA`re-prefill` 段是 KVC TTFT p99 的主体。
2. **目标**D 在 direct-to-D append 完成后异步把新 KV 增量推回 P。reseed 触发时 P 已经有 fresh snapshot → 跳过 model.forward()、直接复用 KV 做 P→D' 传输。
3. **决策**:选 Option C —— **D→P snapshot 按 append-completion 推送P 端用独立 PrefillSnapshotStore 存储(不进 radix treeprefill 在有 snapshot 时 bypass 计算只触发传输**
4. **拒绝的 alternatives**A让 P radix tree 接受多生产者写入§4.3 工程灾难、BD→D' 直推,绕过 P但 mooncake 无 D-Sender 角色 + session-not-resident 场景失败、D仅 eviction 时推async 来不及 + sync 拖死 eviction
5. **工程量**~600 LOC拆 6-8 commit。最难的是 mooncake 双角色化的 thread-safety 和 P 端 prefill bypass 的调度器 hook。
6. **必须 RDMA**:所有传输走 mooncake batch_transfer不允许 TCP fallback。
---
## 1. 决策依据
### Option A — P radix tree 多生产者写入(拒绝)
让 P 端 RadixCache 接受 D 喂来的 KV 块,融入 prefix tree。
**为何拒绝**
- SGLang radix tree 假设单生产者(本 worker 的 model 输出)。改动涉及节点写入路径、引用计数、跨 worker 数据格式、eviction policy 协调。
- 工程量 ~1-2 周,且是侵入式改动,长期维护成本高。
- 与 vendor 上游 diff 太大,未来 rebase 风险高。
### Option B — D→D' 直推(拒绝)
migration 时 D_old 把 KV 直接发到 D_new绕过 P。
**为何拒绝**
- 触发条件 `session-not-resident` 时 KV 已 freeD_old 拿不到任何数据可推。
- mooncake DECODE 模式当前只有 receiver 角色(`assert disaggregation_mode == PREFILL` at conn.py:1563新增 D-Sender 角色与 P-Receiver 角色对偶,工程量与 Option C 相当但**只 cover 部分场景**。
- D→D' 控制平面需要额外协调("哪个 D 当前持有 session"),增加路由复杂度。
### Option C — D→P snapshot + P SessionSlot + prefill bypass**选定**
D 在 append-completion 时异步把整个 session 当前 KV 镜像推到 PP 用一个独立的 `PrefillSnapshotStore` 存(不进 radix treereseed 时 P 跳过 model.forward(),直接用 snapshot 触发 P→D' 传输。
**为何选它**
1. **P 端不动 radix tree**——SnapshotStore 是侧表,无 multi-producer 问题
2. **mooncake 改动局部化**——只放开 `add_transfer_request` 的 PREFILL assertion + 在 DECODE 模式启动一个独立 snapshot transfer 线程
3. **可以分阶段验证**——D→P 推 → P 收到 → P 存 → P 用,每一步可独立 smoke test
4. **failure semantics 干净**——snapshot 缺失就 fallback 到现有 re-prefill 路径,零回退风险
5. **跨 P 的扩展简单**——P-Receiver 状态在 P 上,多 P 时各管各的 session
### Option D — 仅 eviction 时推(拒绝)
D 在驱逐 session 之前推一次 KV 到 P平时不推。
**为何拒绝**
- async 推送reseed 触发时(下一 turn 到达)可能 push 还没到 P 完。需要 reseed path 等 push 完成 → 把延迟成本只是搬家。
- sync 推送:让 eviction 等 mooncake transfer 完,**当前 incoming request触发 eviction 的那个)** 直接被拖死 1-3s。比当前 reseed 还差。
- 不能 cover 非 eviction 触发的 reseed如 migration、admission-no-d-capacity
---
## 2. 架构
```
+---------------- D worker (decode_thread + new snapshot_sender_thread) -----+
| |
| direct-to-D append done |
| | |
| v |
| on_session_step_committed(session_id, kv_committed_len, kv_indices) |
| | |
| v |
| SnapshotSendQueue [throttle by token-delta >= K_DELTA] |
| | |
| v |
| KVSnapshotSender |
| | |
| | mooncake batch_transfer (RDMA) |
| v |
+-----------------------------|----------------------------------------------+
|
v
+---------------- P worker (prefill_thread + new snapshot_receiver_thread) ---+
| |
| KVSnapshotReceiver listening (ZMQ control + mooncake data) |
| | |
| v |
| PrefillSnapshotStore[session_id] -> SnapshotEntry { |
| req_pool_idx, kv_indices, kv_committed_len, last_recv_time |
| } |
| |
| When prefill request arrives with session_id + snapshot_token: |
| | |
| v |
| prefill_bypass_check(session_id, requested_seq_len) |
| | hit: skip model.forward, reuse stored kv, fire P→D' transfer |
| | miss: fall through to normal prefill |
+----------------------------------------------------------------------------+
+--------------- agentic-pd-hybrid (replay.py) -------------------------------+
| |
| _invoke_kvcache_seeded_router (reseed entry): |
| 1. GET /v1/sessions/{sid}/snapshot_status on P → seqlen |
| 2. if seqlen >= requested input_len: |
| set request header x-prefill-use-snapshot=1 |
| route to P → P uses bypass path |
| else: |
| normal seeded_router (re-prefill) |
+----------------------------------------------------------------------------+
```
---
## 3. 数据流时间线
### 3.1 Direct-to-D append + 异步 D→P push
```
t=0 turn N 到 D走 direct-to-D append-prefill
t=T1 direct append 完成scheduler 调 cache_finished_req
SessionAwareCache.cache_finished_req 把 KV 写回 SessionSlot
(此时 KV 全在 D 的 kv_pool 里slot 持锁)
t=T1+ε D-side hook: on_session_step_committed(sid, slot)
计算 delta = slot.kv_committed_len - last_pushed_seqlen[sid]
if delta >= K_DELTA (默认 1024 tokens): 入队 SnapshotSendQueue
t=T1+δ snapshot_sender 线程取出 entry → mooncake batch_transfer
把 kv_pool[slot.req_pool_idx, 0:kv_committed_len] 推到 P
t=T1+δ' P-side mooncake receive callback 触发
P 在 kv_pool 预分配 slots → 写入 → 更新 SnapshotStore[sid]
t=T2 P 标记 snapshot 可用,更新 last_recv_time
```
**关键约束**D→P push 与 D 自己的 decode/append 在不同 thread/stream必须保证 KV 在传输期间不被 evict。
- 复用 SessionSlot 的 lock_ref 机制snapshot_sender 在传输期间 hold lock传输完后 dec_lock。
- 如果 session 在传输期间被 release_session 调用snapshot 应该 abort数据不一致
### 3.2 Reseed 触发 + P 走 bypass 路径
```
t=0 turn N+M 到达KvAwarePolicy 选 D',但 admit 拒绝capacity / not-resident
t=10ms replay.py 进入 _invoke_kvcache_seeded_router
t=15ms probe: GET p/v1/sessions/{sid}/snapshot_status -> {seqlen: 50080, fresh: true}
t=20ms replay: 50080 >= request.input_length (49800),触发 bypass 路径
t=25ms open D' streaming session (HTTP)
t=30ms open P streaming session, set x-prefill-use-snapshot header
t=40ms forward request to SGLang pd-router → P
t=45ms P scheduler 看到 use-snapshot 标记
→ SnapshotStore.lookup(sid) -> SnapshotEntry
→ 跳过 model.forward()
→ 直接复用 SnapshotEntry.kv_indices 给 mooncake KVSender
t=50ms mooncake P→D' RDMA transfer 启动
t=300ms P→D' 完成D' 上 session 重建
t=305ms D' 开始 decode
t=350ms first token 出来 → TTFT
```
**收益对照**
| 段 | 当前 reseed | bypass 后 |
|---|---:|---:|
| P open session | ~50ms | ~50ms |
| **P re-prefill** | **~1500-3000ms** | **0** |
| P→D' transfer (RDMA) | ~200-400ms | ~200-400ms |
| D' decode start | ~50ms | ~50ms |
| TTFT 总 | ~1.8-3.5s | ~0.3-0.5s |
---
## 4. 接口和数据结构
### 4.1 Mooncake 双角色
**Change**: `MooncakeKVManager.__init__` 在 DECODE 模式下**额外**启动 snapshot sender 基础设施(独立 transfer_queues + thread pool
```python
# In MooncakeKVManager.__init__, after start_decode_thread() in DECODE mode:
if envs.SGLANG_DTOP_SNAPSHOT_ENABLED.get():
self._init_snapshot_sender() # new
def _init_snapshot_sender(self):
self.snapshot_send_queue: FastQueue = FastQueue()
self.snapshot_executor = ThreadPoolExecutor(max_workers=2)
threading.Thread(
target=self._snapshot_send_worker,
daemon=True,
).start()
```
**Change**: 删除 `add_transfer_request``assert PREFILL`,改为按 caller 路径分发:
- `add_transfer_request` —— prefill 用,保持现状
- `add_snapshot_transfer_request` —— 新增decode 用
### 4.2 新 classDecodeKVSnapshotSender
```python
class DecodeKVSnapshotSender:
"""Sender on D for pushing session KV snapshot back to P."""
def __init__(self, mgr: MooncakeKVManager, target_p_addr: str,
target_p_bootstrap_room: int, session_id: str):
...
def send(self, kv_indices: npt.NDArray[np.int32],
kv_committed_len: int, aux_blob: bytes) -> None:
"""Enqueue snapshot for async push. Non-blocking."""
def poll(self) -> KVPoll: ...
```
### 4.3 P 端 PrefillSnapshotStore + Receiver
```python
@dataclass
class SnapshotEntry:
session_id: str
req_pool_idx: int
kv_indices: torch.Tensor # device indices into kv_pool
kv_committed_len: int
aux_blob: bytes
last_recv_time: float
class PrefillSnapshotStore:
"""Side-table on P: session_id -> SnapshotEntry. NOT in radix tree."""
def __init__(self, kv_pool_allocator, req_to_token_pool, max_sessions: int = 8):
self.entries: dict[str, SnapshotEntry] = {}
self.max_sessions = max_sessions
...
def ingest(self, session_id: str, kv_data: torch.Tensor,
kv_committed_len: int, aux_blob: bytes) -> None:
"""Allocate slots, copy KV in, register entry. LRU-evicts when full."""
def lookup(self, session_id: str) -> Optional[SnapshotEntry]: ...
def release(self, session_id: str) -> None:
"""Free the slots + remove entry."""
```
### 4.4 P-side prefill bypass 调度器 hook
**Change**: `scheduler.py``handle_generate_request` 入口处检查 `x-prefill-use-snapshot` header / `session_params.use_snapshot=True`
```python
if snapshot_requested and self._snapshot_store.has(session_id):
entry = self._snapshot_store.lookup(session_id)
if entry.kv_committed_len >= len(input_ids) - K_TAIL_TOLERANCE:
return self._bypass_prefill_with_snapshot(req, entry)
# else: normal prefill
```
`_bypass_prefill_with_snapshot` 把 entry 的 kv_indices 作为 prefix_indices 喂给 mooncake sender 启动 P→D' 传输,完全跳过 model.forward()。
### 4.5 D 端 commit hook
**Change**: `scheduler.py``handle_finish_request` / `cache_finished_req` 完成后调用:
```python
if (self._enable_d_to_p_sync and req.session and req.session.streaming
and self._has_p_snapshot_target(req.session.session_id)):
self._maybe_enqueue_snapshot_push(req.session.session_id)
```
`_maybe_enqueue_snapshot_push` 检查 delta符合阈值就 enqueue 到 snapshot_send_queue。
### 4.6 HTTP endpoints (P)
```
GET /v1/sessions/{sid}/snapshot_status
-> {"exists": bool, "seqlen": int, "freshness_s": float}
POST /v1/sessions/{sid}/snapshot_target
-> {"bootstrap_addr": str, "bootstrap_room": int}
(D queries this once per session to learn where to push)
```
### 4.7 agentic-pd-hybrid hook
**File**: `src/agentic_pd_hybrid/replay.py`
In `_invoke_kvcache_seeded_router`, before opening P session:
```python
if config.enable_d_to_p_sync:
snapshot_status = await _probe_p_snapshot(
client, prefill_url, session_id, target_seqlen=request.input_length,
)
if snapshot_status and snapshot_status["fresh"]:
# bypass path
return await _invoke_kvcache_snapshot_bypass(...)
# else: existing seeded router
```
### 4.8 CLI flag
```
--enable-d-to-p-sync (default off)
--d-to-p-sync-delta-tokens (default 1024)
--d-to-p-sync-max-sessions (default 8 on P)
```
---
## 5. 实现路线图(每步独立 commit
| # | Commit subject | Files | Why a separate commit |
|---|---|---|---|
| 1 | `feat(sglang): mooncake bidirectional infra for D→P snapshot` | `third_party/sglang/.../mooncake/conn.py` | 隔离 mooncake 层改动;不破坏 PD-disagg 现有路径 |
| 2 | `feat(sglang): PrefillSnapshotStore + DecodeKVSnapshotSender` | `third_party/sglang/.../mem_cache/`, `third_party/sglang/.../disaggregation/mooncake/` | 新数据结构 |
| 3 | `feat(sglang): P-side prefill bypass with snapshot` | `third_party/sglang/.../managers/scheduler.py`, `tokenizer_manager.py` | 调度器 hook最危险单独提交便于回滚 |
| 4 | `feat(sglang): D-side session commit hook → snapshot push` | `third_party/sglang/.../managers/scheduler.py`, `session_aware_cache.py` | D 端 trigger |
| 5 | `feat(sglang): HTTP endpoints for snapshot status/target` | `third_party/sglang/.../entrypoints/http_server.py` | API 表面 |
| 6 | `feat(agentic): D→P sync hook in seeded_router` | `src/agentic_pd_hybrid/replay.py` | 客户端逻辑 |
| 7 | `feat(agentic): --enable-d-to-p-sync CLI + config` | `src/agentic_pd_hybrid/cli.py`, `benchmark.py` | CLI 接入 |
| 8 | `feat(experiments): smoke test + E4 sweep scripts` | `scripts/`, `docs/D_TO_P_SMOKE_RESULTS_ZH.md` | 验收 + 落盘 |
---
## 6. Metrics + 观察性
### Structural log channels写到 `structural/d-to-p-sync.jsonl`
```json
{"ts": ..., "event": "snapshot_push_enqueued", "sid": "...", "delta": 2048}
{"ts": ..., "event": "snapshot_push_sent", "sid": "...", "bytes": 4_200_000_000, "dur_ms": 320}
{"ts": ..., "event": "snapshot_push_failed", "sid": "...", "reason": "..."}
{"ts": ..., "event": "snapshot_recv_ingested", "sid": "...", "seqlen": 50000}
{"ts": ..., "event": "snapshot_evicted", "sid": "...", "reason": "lru|session_close|stale"}
{"ts": ..., "event": "snapshot_bypass_hit", "sid": "...", "seqlen": 50000, "saved_prefill_ms_est": 1800}
{"ts": ..., "event": "snapshot_bypass_miss", "sid": "...", "reason": "no_entry|stale|seqlen_short"}
```
### Per-request metrics (additional fields in metrics.jsonl)
```
d_to_p_snapshot_used: bool
d_to_p_snapshot_age_s: float | None
d_to_p_push_count_during_session: int
```
### Sweep summary 应回答的问题
1. snapshot push 触发频率(每秒多少次)
2. snapshot LRU eviction 是不是瓶颈freshness 分布)
3. reseed 触发时 bypass hit rate
4. bypass vs fallback 的 TTFT 分布对比
---
## 7. 失败模式 + 回退
| 失败模式 | 现象 | 处理 |
|---|---|---|
| D→P transfer 中途失败 | mooncake KVPoll.Failed | snapshot_send_queue 重试 1 次,再失败放弃;保留旧 entry |
| P snapshot store 满 | LRU 淘汰最旧 entry | log eviction event |
| reseed 时 snapshot stale | entry.kv_committed_len < requested input_len - K_TAIL_TOLERANCE | 回退到 normal re-prefill |
| D 重启 / session 丢失 | D session_aware_cache 没了 | snapshot_target 注册过期下次 push 收到 404 清理 D 端记录 |
| P 重启 | snapshot store 清空 | 下次 reseed probe 拿到 not-exists fallback |
| 双重 push多个 D 喂同一 session| 不该发生session 同时只在一个 D但保险起见用 last-write-wins + log warning | |
**核心不变量**DP sync 失败永远只导致 fallback 到现有 re-prefill 路径不影响正确性
---
## 8. 测试
### Smoke test 阶段commit #8
`scripts/smoke_d_to_p_sync.sh`
1. 1P1D开启 `--enable-d-to-p-sync`
2. 5 sessions × 3 turns 的迷你 trace
3. 触发条件第二 turn direct-to-D append 完成后强制 capacity-evict admission flag 调小
4. 第三 turn 必然走 reseed 路径
5. 验证
- structural log snapshot_push_sent + snapshot_recv_ingested
- 第三 turn metrics 显示 d_to_p_snapshot_used=true
- TTFT cold prefill 的差异 1s
### E4 端到端 sweepfeature 验收完成后)
详见 §9
---
## 9. 实验E4 KVC w/ D→P vs naive PD-disagg
**目标**证明 KVC + DP 在保持 session affinity 设计独特性的前提下 latency 优于 naive PD-disaggE1 baseline)。
### 实验矩阵
| # | 配置 | 期望验证 |
|---|---|---|
| E1已有 | naive 1P3D + kv-aware + RDMA | baseline KVC |
| E3已有 | KVC v2 + RDMA + load-floor | KVC 但无 DPreseed prefill |
| **E4** | KVC v2 + RDMA + load-floor + DP | KVC + DP bypass |
| E4-ablate | KVC v2 + RDMA + load-floor + DP但人为 disable bypass | 排除 push 流量本身的副作用 |
### 假设
- **H4-1**E4 TTFT p99 E1证明KVC + DP p99 长尾上不再输 naive PD-disagg
- **H4-2**E4 reseed 占比execution_mode=*reseed*)不变,但 reseed 路径自身 TTFT 中位 E1 normal 路径 TTFT 中位
- **H4-3**E4 的总 throughput 略低于 E3因为 DP 推送占带宽 TTFT/latency 优势足以补偿
### 数据集
- `outputs/inferact_50sess.jsonl` E1/E2/E3
- md5 7bb263a32600ef5a6ef5099ba340a487
### 报告(事前 commit `docs/E4_PROTOCOL_ZH.md`,跑完后 `docs/E4_RESULTS_ZH.md`
每个 hypothesis 标注
- 证实 / 证伪 / 部分证实
- 数字证据
- 失败原因若证伪
- 后续工作建议
---
## 10. 边界 + 非目标
**本设计不解决**
- **DD' 直推**未来若证实场景 X 必须用可走 Option B 作为补充
- ** P 协调**现假设单 P P 时每个 P 各自维护自己的 snapshot storesession 路由到哪个 P router 决定
- **跨节点 mooncake**当前 H200 是单机 4 GPUIB device mlx5_60跨节点 RDMA 留作 future work
- **snapshot 持久化**P 重启 snapshot 全丢下次 reseed fallback不写盘
- **prefill bypass chunked prefill 的交互**bypass 走的是 " session KV 直接传输"不和 chunked prefill 并存 P 当前正在 chunked-prefill 这个 sessionbypass 等到现有 chunk 结束再起
---
## 11. 决策点(等评审)
| # | 问题 | 默认 |
|---|---|---|
| D1 | snapshot push throttle delta K_DELTA = 1024 tokens 合理太小会泛滥推送太大会让 snapshot 滞后 | 起步用 1024 smoke 看流量再调 |
| D2 | snapshot LRU 上限 max_sessions = 8 合理P ~92K tokenssession 平均 50K 1-2 | 8 太乐观 4 |
| D3 | bypass P 是否走 mooncake staging buffer还是直接 zerocopy | 直接 zerocopy避免一次 devicedevice 拷贝 |
| D4 | D-side push 失败后是否上报 router 影响策略 | 不上报fail-openfallback re-prefill 也能跑 |
| D5 | snapshot 是否包含 aux/statemamba state, swa 状态等 | E4 实验 trace 只用 Qwen3 mambaaux 跟着 KV 一起带 |
---
**核心句**DP 同步是 KVC 设计真正击败 naive PD-disagg 的关键缺口本设计用 P 端独立 snapshot store + prefill bypass 的最小改动方案避开 radix tree 多生产者扩展的工程陷阱~600 LOC 8 commit 可在单次 session 完成验收后即可启动 E4 实验对比 KVC vs naive