# 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 tree），prefill 在有 snapshot 时 bypass 计算只触发传输**。 4. **拒绝的 alternatives**：A（让 P radix tree 接受多生产者写入，§4.3 工程灾难）、B（D→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 已 free，D_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 镜像推到 P；P 用一个独立的 `PrefillSnapshotStore` 存（不进 radix tree）；reseed 时 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 新 class：DecodeKVSnapshotSender ```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 | | **核心不变量**：D→P 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 端到端 sweep（feature 验收完成后） 详见 §9。 --- ## 9. 实验：E4 KVC w/ D→P vs naive PD-disagg **目标**：证明 KVC + D→P 在保持 session affinity 设计独特性的前提下 latency 优于 naive PD-disagg（E1 baseline）。 ### 实验矩阵 | # | 配置 | 期望验证 | |---|---|---| | E1（已有） | naive 1P3D + kv-aware + RDMA | baseline，无 KVC 层 | | E3（已有） | KVC v2 + RDMA + load-floor | KVC 但无 D→P，reseed 重 prefill | | **E4** | KVC v2 + RDMA + load-floor + D→P | KVC + D→P bypass | | E4-ablate | KVC v2 + RDMA + load-floor + D→P，但人为 disable bypass | 排除 push 流量本身的副作用 | ### 假设 - **H4-1**：E4 的 TTFT p99 ≤ E1。证明：KVC + D→P 在 p99 长尾上不再输 naive PD-disagg。 - **H4-2**：E4 的 reseed 占比（execution_mode=*reseed*）不变，但 reseed 路径自身 TTFT 中位 ≤ E1 normal 路径 TTFT 中位。 - **H4-3**：E4 的总 throughput 略低于 E3（因为 D→P 推送占带宽），但 TTFT/latency 优势足以补偿。 ### 数据集 - `outputs/inferact_50sess.jsonl`（同 E1/E2/E3） - md5 7bb263a32600ef5a6ef5099ba340a487 ### 报告（事前 commit `docs/E4_PROTOCOL_ZH.md`，跑完后 `docs/E4_RESULTS_ZH.md`） 每个 hypothesis 标注： - 证实 / 证伪 / 部分证实 - 数字证据 - 失败原因（若证伪） - 后续工作建议 --- ## 10. 边界 + 非目标 **本设计不解决**： - **D→D' 直推**：未来若证实场景 X 必须用，可走 Option B 作为补充 - **跨 P 协调**：现假设单 P。多 P 时每个 P 各自维护自己的 snapshot store，session 路由到哪个 P 是 router 决定 - **跨节点 mooncake**：当前 H200 是单机 4 GPU，IB device 选 mlx5_60。跨节点 RDMA 留作 future work - **snapshot 持久化**：P 重启 snapshot 全丢，下次 reseed 走 fallback。不写盘 - **prefill bypass 与 chunked prefill 的交互**：bypass 走的是 "全 session KV 直接传输"，不和 chunked prefill 并存。若 P 当前正在 chunked-prefill 这个 session，bypass 等到现有 chunk 结束再起 --- ## 11. 决策点（等评审） | # | 问题 | 默认 | |---|---|---| | D1 | snapshot push 的 throttle delta K_DELTA = 1024 tokens 合理？太小会泛滥推送，太大会让 snapshot 滞后 | 起步用 1024，跑 smoke 看流量再调 | | D2 | snapshot LRU 上限 max_sessions = 8 合理？P 池 ~92K tokens，session 平均 50K → 1-2 个？ | 8 太乐观，改 4 | | D3 | bypass 时 P 是否走 mooncake 的 staging buffer？还是直接 zerocopy | 直接 zerocopy，避免一次 device→device 拷贝 | | D4 | D-side push 失败后是否上报 router 影响策略？ | 不上报，fail-open（fallback re-prefill 也能跑） | | D5 | snapshot 是否包含 aux/state？（mamba state, swa 状态等） | E4 实验 trace 只用 Qwen3，无 mamba。aux 跟着 KV 一起带 | --- **核心句**：D→P 同步是 KVC 设计真正击败 naive PD-disagg 的关键缺口。本设计用 P 端独立 snapshot store + prefill bypass 的最小改动方案，避开 radix tree 多生产者扩展的工程陷阱，~600 LOC 拆 8 commit 可在单次 session 完成。验收后即可启动 E4 实验对比 KVC vs naive。