agentic-pd-hybrid/docs/SNAPSHOT_STORE_REFACTOR_ZH.md

SnapshotStore 重构（解决 P-side alloc-failed 死局）

日期：2026-05-13 Status：设计阶段，开始实施 根因：docs/E4_VS_E1_RESULTS_ZH.md §3 + E4-v4/v5 forensic 显示 D→P sync 167 次尝试 0 OK，全部因 prepare_receive 试图从 token_to_kv_pool_allocator.alloc(N) 拿 N 个 slot 而 P 的池被自己 prefill 工作占满

---

0. TL;DR

• 当前 P-side prepare_receive 用 token_to_kv_pool_allocator.alloc(N) 抢 kv_pool slot —— 跟 P 自己的 prefill 工作直接争抢资源 → 90%+ 时间 alloc-failed
• 重构方向：P-side 用独立 GPU buffer 接收 snapshot，与 kv_pool 解耦
• 在 finalize_ingest 时才把 snapshot bytes copy 进 kv_pool slots（此时可以等更优的时机）
• ~250 LOC 新代码，主要在 disaggregation/snapshot/controller.py

---

1. 当前实现的死局

prepare_receive(sid, num_tokens=50000):
    indices = self.token_to_kv_pool_allocator.alloc(50000)
    if indices is None:
        return ok=False, reason="alloc-failed"   ← 90%+ 时间走这里
    return slot_indices = indices.tolist()

alloc(50000) 在 P 池中找 50000 个 contiguous 空 slot。当 P 正在 prefill 自己的 request 时（这是 P 的常态），池里大部分 slot 被锁定 → 找不出 50K 个空闲的 → fail.

E4-v5 167 次 sync 尝试统计：

• 148 个 alloc-failed（88%）
• 19 个 session-not-resident（D 端已 evict）
• 0 个 OK

---

2. 新设计：PrefillSnapshotStore 侧表

   ┌─────────────────────────────────────────────────────────────────┐
   │ P worker scheduler                                               │
   │                                                                  │
   │  kv_pool (existing, owned by P's prefill work)                  │
   │  ┌────────────────────────────────────────────────┐             │
   │  │ k_buffer[0..L]: (max_tokens, head, dim)        │             │
   │  │ v_buffer[0..L]: (max_tokens, head, dim)        │             │
   │  └────────────────────────────────────────────────┘             │
   │                                                                  │
   │  snapshot_buf (NEW, dedicated for D→P snapshot reception)       │
   │  ┌────────────────────────────────────────────────┐             │
   │  │ pinned GPU tensor of size SNAPSHOT_BUF_BYTES   │             │
   │  │ (default 8 GB)                                  │             │
   │  │ • registered with mooncake (one-time at init)  │             │
   │  │ • slab-allocator manages free space             │             │
   │  └────────────────────────────────────────────────┘             │
   └─────────────────────────────────────────────────────────────────┘

Flow:
  1. prepare_receive(sid, N):
       slab = snapshot_buf_allocator.alloc(N * per_token_bytes_total)
       record = (sid, slab_offset, N)
       return (snapshot_buf_base + slab_offset for K_L, V_L per layer)
       ← never blocks on kv_pool

  2. (out-of-band) D pushes KV bytes into the slab via mooncake RDMA

  3. finalize_ingest(sid, token_ids):
       record = pop ingest_record[sid]
       slots = token_to_kv_pool_allocator.alloc(N)  ← can fail here
       if alloc-failed:
           snapshot_buf_allocator.free(record.slab)
           return ok=False, reason=alloc-failed-on-finalize
       # copy snapshot_buf[layer L][token range] → kv_pool.k_buffer[L][slots]
       for L in range(layer_num):
           kv_pool.k_buffer[L][slots] = snapshot_buf[K_L_offset : K_L_offset + N * K_stride].view(N, head, dim)
           kv_pool.v_buffer[L][slots] = snapshot_buf[V_L_offset : V_L_offset + N * V_stride].view(N, head, dim)
       tree_cache.insert(InsertParams(key=token_ids, value=slots))
       snapshot_buf_allocator.free(record.slab)
       return ok=True

---

3. 关键 design choices

决策	选择	原因
Snapshot buffer 存哪	GPU memory	与 D RDMA 目标对称（D 端 KV 也在 GPU），避免 host↔device 拷贝
默认大小	8 GB	Qwen3-30B 一个 ~50K-token session 的 KV ~5 GB；8 GB 让我们至少 hold 一个 + 部分备份
分配粒度	单次 contiguous 一个 session 全部 KV	简化 slab allocator + 单次 batch transfer
Layout	K-all-layers concat, then V-all-layers concat	跟 mooncake 的 batch_transfer 接口对齐
Free 策略	finalize 后立即 free	当 snapshot 已 ingest 到 kv_pool，snapshot_buf 副本不再需要
满了怎么办	prepare_receive 返回 ok=False, reason=snapshot-buf-full	让 caller fall back 到 re-prefill

---

4. 接口变化

4.1 SnapshotPrepareReceiveReqOutput

旧：

k_base_ptrs: List[int]   # 各 layer 的 k_buffer.data_ptr()
v_base_ptrs: List[int]
slot_indices: List[int]  # kv_pool 中分配的 slot
stride_k_bytes / stride_v_bytes

新：

snapshot_buf_base_ptr: int  # snapshot_buf.data_ptr()
k_layer_offsets: List[int]  # 各 layer K 在 snapshot_buf 中的字节偏移
v_layer_offsets: List[int]  # 各 layer V 偏移
num_tokens: int
stride_k_bytes / stride_v_bytes
slab_handle: int            # opaque handle for finalize/abort

4.2 SnapshotFinalizeIngestReqInput

旧：

session_id, token_ids, slot_indices

新：

session_id, token_ids, slab_handle   # P 用 handle 找到 record，再 alloc kv_pool + copy + insert

4.3 D-side push 逻辑（agentic）

旧：D 算 src_slot[L] → dst_slot[L] mapping，batch_transfer

新：D 算 src_slot[L] → snapshot_buf 中的 k_layer_offsets[L] / v_layer_offsets[L] mapping，batch_transfer。完全不需要 dst slot indices。

---

5. 实施步骤

#	步骤	LOC 估计
1	SnapshotBufAllocator 类（slab/bump allocator）	80
2	SnapshotLinkController.__init__ 加 snapshot_buf 分配 + 注册	30
3	重写 prepare_receive、新加 _compute_layer_offsets	60
4	新加 finalize_with_snapshot_buf + 删旧的 finalize_ingest	70
5	修改 io_struct 字段 + 删旧字段	30
6	修改 agentic _attempt_d_to_p_sync 用新字段	40
7	改 mem leak check 计入 snapshot_buf	5
8	单元 smoke test	50

Total: ~365 LOC

---

6. 风险

风险	缓解
8 GB GPU mem cost	用户可配置；mem-fraction-static 已经留了 buffer
多 session 抢 snapshot_buf	slab allocator + LRU evict 旧的 snapshot
GPU→GPU copy 性能	~5 GB @ 3 TB/s = 1.7 ms，可忽略
接口大改影响 smoke	在 commit 内完成所有接口变更，smoke 同步更新

---

7. 验收

• scripts/smoke_snapshot_sglang_integration.py 跑通新接口（prepare_receive 不再 alloc-failed）
• E4-v6 跑同样 trace，d-to-p-sync.jsonl 出现 OK 事件 ≥ 30%（vs 当前 0%）

---

核心句：用 GPU 上独立的 snapshot_buf 接收 D 端推送，把"竞争 P kv_pool"这个根本性 alloc 冲突消掉，把 alloc 决策推迟到 finalize 时机，让 D→P 真正有机会跑通。