docs(design): block-level eviction refactor — concrete API plan

Turns the architectural manifesto (KVC_EVICTION_GRANULARITY_DESIGN_ZH.md) into a function-by-function design the next collaborator can implement against. Contents: - §1 current SessionAwareCache state with exact field semantics (req_pool_idx / kv_committed_len / kv_allocated_len / cache_protected_len) - §3.1–§3.6 post-refactor source sketches for SessionSlot, cache_finished_req, cache_unfinished_req, match_prefix, release_session, get_session_status - §3.7 the schedule_batch.py:1572-1646 correction block we can remove (the E3 landmine) - §4 five invariants the PR must defend - §5 GPU-free unit + property test plan with a MockRadixCache shape - §6 ~1 week engineering estimate and three risks - §7 dependency relationship to the planned D->P sync work - §8 minimal step list for the implementing agent No code change yet. Future commits on a feat/block-level-evict branch will execute against this spec.
2026-05-12 23:49:18 +08:00
parent baa843a3f9
commit 683c44bd71
1 changed files with 309 additions and 0 deletions
--- a/docs/BLOCK_LEVEL_EVICTION_DESIGN_ZH.md
+++ b/docs/BLOCK_LEVEL_EVICTION_DESIGN_ZH.md
@@ -0,0 +1,309 @@
+# Block-level Eviction Refactor — 设计文档
+
+**日期**：2026-05-12
+**前置**：[KVC_EVICTION_GRANULARITY_DESIGN_ZH.md](KVC_EVICTION_GRANULARITY_DESIGN_ZH.md)（架构层 manifesto）
+**性质**：实现层设计 + API 草案 + 测试计划，供下一个合作者直接据此编码
+**Status**：草案，未实现。代码全部 quoted from `third_party/sglang/python/sglang/srt/mem_cache/session_aware_cache.py @ origin/h200-cu130`
+
+---
+
+## 0. TL;DR
+
+把 `SessionAwareCache` 当前对 streaming-session **整段 KV 一次性 free** 的语义改成：
+
+1. Streaming-session decode 输出在 turn finish 时 **增量 commit 进 radix tree**。
+2. `SessionSlot` 退化为**纯 metadata**（仅持 `last_node` + lock_ref 状态），不再独占 KV 区间。
+3. `release_session` 改为只 dec_lock_ref + 删 slot，**让 SGLang 标准 radix LRU 按 block 粒度蚕食**。
+
+预期收益：evict 粒度从一次 ~67K tokens 降到 ~24 tokens（page_size 个 token），reseed 频率降一个数量级；同时把 P 端 radix tree 改造成可被外部喂数据（为 [D_TO_P_SYNC_CONTRACT_ZH.md](D_TO_P_SYNC_CONTRACT_ZH.md) 铺路）。
+
+---
+
+## 1. 现状代码梳理
+
+### 1.1 关键文件与函数
+
+`third_party/sglang/python/sglang/srt/mem_cache/session_aware_cache.py`
+
+| 函数 / 字段 | 当前语义 |
+|---|---|
+| `SessionSlot.req_pool_idx` | streaming-session 独占的 req_pool 槽位 |
+| `SessionSlot.kv_committed_len` | 上一 turn 完成时已 commit 的 KV 长度（已计入 cache_protected_len 部分进入 radix） |
+| `SessionSlot.kv_allocated_len` | 当前已分配但**未进 radix** 的 KV 长度（"session-exclusive 尾部"） |
+| `SessionSlot.cache_protected_len` | 首 turn 提交 radix 时的 protected 边界 |
+| `match_prefix(streaming req)` | 命中 slot → 返回 `req_to_token[req_pool_idx, :prefix_len]`，bypass radix |
+| `cache_unfinished_req(streaming req)` | subsequent turns → **完全 skip inner**（不进 radix） |
+| `cache_finished_req(streaming req)` | 调 `slot.save_from_req`，**不调 inner.cache_finished_req** |
+| `release_session(sid)` | `dec_lock_ref(slot.last_node)` + `free(req_to_token[req_pool_idx, cache_protected_len:kv_allocated_len])` + 回收 req_pool 槽位 |
+
+### 1.2 当前为什么是错的（重述）
+
+`[cache_protected_len, kv_allocated_len)` 是首轮入 radix 之后所有累积的 decode 输出 + 后续 turn 的 extend。在 Inferact / SWE-Bench 实测：
+
+- `cache_protected_len` ≈ 首 turn boilerplate ~12K
+- `kv_allocated_len` 累积 50–100K
+- 每次 `release_session` 一次性释放 38–88K，这部分**从未进 radix**，无法享受 leaf-by-leaf 渐进 evict
+
+→ session 被 evict 后必须从 client 原 prompt 重 prefill 全长 + mooncake transfer 全长，跟 naive PD-disagg 等价（详见 manifesto §1）。
+
+---
+
+## 2. 目标行为表
+
+| 场景 | 现状 | 目标 |
+|---|---|---|
+| Session 累积 50K KV，D 满了 | `release_session` 一次释放 38K | radix LRU 从最老 leaf 开始 evict，单次 ~24 tokens |
+| Session 被 evict 后再到来 | 必须 reseed 50K | 仅 re-prefill 被 evict 的 leaf 部分（典型 ≤ 5K） |
+| Evicted session TTFT | 50–90K reseed ≈ 3–7s | 5K append-prefill ≈ 200ms |
+| 不被 evict 的 session | 同 session 内 turns append-only | 同样 append-only（不变） |
+| Direct-to-D fast path 命中率 | 91.6% (SWE-Bench) / 38% (E3 Inferact) | 应 ≥ 85% 即使 saturation |
+
+---
+
+## 3. 设计
+
+### 3.1 SessionSlot 字段精简
+
+**after refactor**：
+
+```python
+@dataclass
+class SessionSlot:
+    virtual_node: _VirtualNode = field(default_factory=_VirtualNode)
+
+    # Pointer into the radix tree — the deepest node owned by this session's
+    # committed prefix. Held under inc_lock_ref so radix LRU never evicts this
+    # *active* leaf out from under a turn-in-progress. Released by
+    # release_session.
+    last_node: Any = None
+    swa_uuid_for_lock: Optional[str] = None
+
+    # Bookkeeping fields (no longer authoritative ownership of KV indices).
+    last_access_time: float = field(default_factory=time.monotonic)
+
+    # Mamba state stays slot-owned (mamba doesn't fit the radix model).
+    mamba_pool_idx: Any = None
+    mamba_ping_pong_track_buffer: Any = None
+    mamba_next_track_idx: Any = None
+    mamba_last_track_seqlen: Any = None
+    mamba_branching_seqlen: Any = None
+```
+
+**删除**：`req_pool_idx`、`kv_committed_len`、`kv_allocated_len`、`cache_protected_len`、`swa_evicted_seqlen`。这些字段的真值改由 radix tree + req_to_token_pool 共同维护。
+
+### 3.2 `cache_finished_req` 改造
+
+**after refactor**：
+
+```python
+def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
+    if not _is_streaming(req):
+        return self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
+
+    session_id = req.session.session_id
+    slot = self.slots.setdefault(session_id, SessionSlot())
+
+    # KEY CHANGE: always delegate to inner — this inserts the new tokens
+    # (kv_committed_len .. fill_ids end) as radix-tree blocks. Subsequent
+    # match_prefix calls for this session will hit the radix tree directly.
+    result = self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
+
+    # Update slot bookkeeping only (no KV ownership).
+    slot.last_node = req.last_node
+    slot.swa_uuid_for_lock = req.swa_uuid_for_lock
+    slot.last_access_time = time.monotonic()
+
+    # Mamba state still goes through slot.
+    slot.mamba_pool_idx = req.mamba_pool_idx
+    ...
+    return result
+```
+
+**不变量**：
+- `inner.cache_finished_req` 会把 `[kv_committed_len_old, kv_committed_len_new)` 范围内对齐到 page_size 的 KV 插入 radix。这个语义来自 SGLang 标准实现，无需改 inner。
+- `slot.last_node` 现在指向**当前 session 已 commit prefix 的尾节点**，每个 turn 后向前推进。
+- `dec_lock_ref(old_last_node)` + `inc_lock_ref(new_last_node)` 必须在 turn 切换时执行。
+
+### 3.3 `cache_unfinished_req` 改造
+
+streaming session 的 subsequent turn **不再 skip inner**。原因：现在 `match_prefix` 走 radix，chunked-prefill 中间状态也需要 inner 维护：
+
+```python
+def cache_unfinished_req(self, req: Req, **kwargs):
+    if _is_streaming(req) and kwargs.get("chunked", False):
+        # Chunked prefill: forward to inner so the per-chunk extend gets
+        # tracked in the radix LRU access timestamps.
+        ...
+    self.inner.cache_unfinished_req(req, **kwargs)
+```
+
+具体的 chunked 处理细节需要保留对 `prefix_indices` 重建的逻辑（参考当前实现 lines 215–225），但调用 `inner.cache_unfinished_req` 不能 skip。
+
+### 3.4 `match_prefix` 改造
+
+退化为**纯 inner 转发**——SessionSlot 不再持 KV 指针：
+
+```python
+def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+    # No more slot-fast-path. Streaming sessions reuse KV via radix tree
+    # match like every other request.
+    return self.inner.match_prefix(params)
+```
+
+调用方需要的 "这个 session 的 committed prefix 长度" 信息改为通过 `inner.match_prefix(...).device_indices.shape[0]` 推导。
+
+### 3.5 `release_session` 改造
+
+**after refactor**：
+
+```python
+def release_session(self, session_id: str) -> int:
+    slot = self.slots.pop(session_id, None)
+    if slot is None:
+        return 0
+
+    # Just release our radix lock — radix LRU can now reclaim our prefix
+    # leaves at its own pace. NO direct token_to_kv_pool free.
+    if slot.last_node is not None:
+        if slot.swa_uuid_for_lock is not None:
+            self.inner.dec_lock_ref(
+                slot.last_node,
+                DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock),
+            )
+        else:
+            self.inner.dec_lock_ref(slot.last_node)
+
+    # Mamba state still needs explicit cleanup if present.
+    if slot.mamba_pool_idx is not None:
+        ...
+
+    return 0  # "freed_tokens" no longer meaningful; radix LRU shed lazily
+```
+
+### 3.6 `get_session_status` / `list_session_statuses` 改造
+
+`resident_tokens` 现在的真值来自 radix tree。需要在 inner 暴露一个 helper：
+
+```python
+# In BasePrefixCache / RadixCache:
+def tokens_under(self, node) -> int:
+    """Count tokens in the path from root to `node` (inclusive)."""
+    ...
+
+# In SessionAwareCache:
+def get_session_status(self, session_id: str) -> Optional[Dict[str, Any]]:
+    slot = self.slots.get(session_id)
+    if slot is None:
+        return None
+    resident_tokens = self.inner.tokens_under(slot.last_node) if slot.last_node else 0
+    return {
+        "session_id": session_id,
+        "resident": resident_tokens > 0,
+        "resident_tokens": int(resident_tokens),
+        "last_access_time": float(slot.last_access_time),
+    }
+```
+
+`admit_direct_append` 的容量检查改用 `resident_tokens` 的 radix 真值（去掉 `kv_committed_len / kv_allocated_len` 双值不一致的可能）。
+
+### 3.7 SGLang 调度路径配套改动
+
+参考 `schedule_batch.py:1572-1646`，当前 streaming-session correction（commit b8e6f13 / 986f351 引入）建立在 SessionSlot 拥有独立 KV 范围之上。block-level refactor 后这条 correction 路径**完全无需存在**——req 的 fill_ids / prefix_indices 由 inner radix `match_prefix` 直接给出一致值。
+
+**移除项**：
+- `schedule_batch.py:1572-1585` 的 `actual_extend_len = max(0, len(fill_ids) - len(prefix_indices))` correction 块。
+- `schedule_batch.py:1646` 的 `assert seq_len - pre_len == req.extend_input_len`（refactor 后该不变量结构上必然成立）。
+- E3 触发的 latent landmine ([E3_FINDINGS_ZH.md](E3_FINDINGS_ZH.md) §2)随之消失。
+
+---
+
+## 4. 不变量（必须在 PR 自测中覆盖）
+
+| Inv | 内容 |
+|---|---|
+| I1 | `release_session(sid)` 后，下一次同 session 请求的 `match_prefix` 行为只取决于 radix tree 的常驻状态——不依赖 `slots` dict。 |
+| I2 | 任意 (session_id, turn_id) 的 `cache_finished_req` 调用后，radix tree 上必然存在一条 root→leaf 路径覆盖该 turn 的全部 committed token（即 `tokens_under(slot.last_node)` 严格不降）。 |
+| I3 | `restore_to_req` 必须**幂等**：在 chunked-prefill 重试场景下，对同一 req 可被调用多次而最终 req 状态等价。当前实现靠"不清 slot 字段"实现 → refactor 后改由 radix `match_prefix` 的纯函数性质保证。 |
+| I4 | 无 streaming-session 的请求（`req.session is None`）行为 **不变**：所有路径 short-circuit 到 inner。 |
+| I5 | 任一 turn 结束后，对 `slot.last_node` 的 `inc_lock_ref` 必须有对应的 `dec_lock_ref`，且 `release_session` 是最终的释放点。 |
+
+---
+
+## 5. 测试计划（无 GPU 可跑）
+
+### 5.1 单元测试（mock inner cache）
+
+写一个 `MockRadixCache(BasePrefixCache)`，记录所有 `cache_finished_req / cache_unfinished_req / match_prefix / evict / dec_lock_ref` 调用序列。然后：
+
+| Test | 断言 |
+|---|---|
+| `test_release_session_no_direct_free` | 调 `release_session` 后，Mock 上 **没有** 直接 `free(kv_indices)` 调用，只有 `dec_lock_ref` |
+| `test_subsequent_turn_inserts_radix` | 模拟 turn 0 → 1 → 2 三次 `cache_finished_req`，断言每次都触发 `inner.cache_finished_req` |
+| `test_match_prefix_uses_inner` | streaming 与 non-streaming 都仅走 `inner.match_prefix` |
+| `test_restore_idempotent` | 模拟 chunked-prefill 重试，连续两次 `match_prefix` 返回的 `device_indices` 一致 |
+| `test_eviction_under_pressure_is_block_level` | inject 一个 "pool 满，必须 evict 24 tokens" 的状态，断言 `release_session` 不被触发，inner 的 LRU 单步走 |
+
+### 5.2 Property-based 测试
+
+```python
+@given(turns=lists(integers(min_value=24, max_value=2048), min_size=1, max_size=50))
+def test_committed_tokens_monotone(turns):
+    """tokens_under(slot.last_node) is monotonically non-decreasing across turns."""
+    ...
+```
+
+### 5.3 Integration smoke（需要 GPU，但放在 sweep 脚本里）
+
+执行 `sweep_e2_kvc_v2_rdma.sh` 同 trace 同配置，对比指标：
+- evict 总次数（期望从 90 → < 10）
+- 单次平均 evict tokens（期望从 67K → < 500）
+- TTFT p99（期望从 1.28s → < 0.7s）
+- direct-to-D 命中率（期望 ≥ 85%）
+
+---
+
+## 6. 工程量与风险
+
+### 6.1 工程量
+
+| 工作 | 估时 | 风险 |
+|---|---|---|
+| §3.1–§3.6 SessionAwareCache 改造 | 2–3 天 | 中：需要熟悉 radix 内部 lock_ref / evict 协议 |
+| §3.7 schedule_batch 清理 | 0.5 天 | 低：是删代码 |
+| §4 不变量单元测试 | 2 天 | 低 |
+| §5.3 GPU smoke + 数据对比 | 2 天 | 中：mooncake 仍可能触发 E2 级联 death，需要 §S3 修复一并跑 |
+| **总计** | **~1 周** | |
+
+### 6.2 关键风险
+
+1. **`inner.cache_finished_req` 对 streaming-session req 的兼容性**：当前 SGLang 标准 radix 假设 req 在 cache_finished_req 时是 "完整 prefill+decode 完成"。streaming-session 的 req 在每个 turn 结束时还会留下"未完成的 conversation"，要确保 inner 在插入时不会把 decode-only tokens 当成可丢弃尾巴。需要 audit `radix_cache.py:cache_finished_req` 的实现。
+
+2. **lock_ref 顺序**：turn N+1 开始的 `match_prefix` → inc_lock_ref(new_node)，turn N 结束的 dec_lock_ref(old_node)，时序若反了会在并发下让 LRU 把刚 commit 的 leaf 误 evict。建议加 assertion：`dec_lock_ref` 之前 `inc_lock_ref` 必须先到。
+
+3. **chunked-prefill retry**：见 I3。SGLang 当前 `restore_to_req` 不清 slot 字段就是为此 retry。refactor 后必须确认 inner radix `match_prefix` 在 retry 下也幂等（标准 radix tree 是的，但要写测试明确锁住这个性质）。
+
+---
+
+## 7. 与 D→P sync 工作的关系
+
+block-level evict 是 [D_TO_P_SYNC_CONTRACT_ZH.md](D_TO_P_SYNC_CONTRACT_ZH.md) 的**前置条件**：
+
+- D→P sync 需要 P 端 radix tree **可接收外部喂入的 KV block**。
+- 当前 P 端 radix 假设单生产者（本 worker 模型输出）。
+- block-level refactor 完成后，streaming-session 的 KV 已经走标准 radix 路径——再让 radix tree 接受"外部喂入"的额外生产者就只是扩展 insert API，而不是发明新的存储路径。
+
+→ 两件事可顺序做：先 block-level evict，再 D→P sync。
+
+---
+
+## 8. 接班 agent 的最小动作
+
+1. fork 一个 `feat/block-level-evict` 分支（从 `improve/audit-and-foundations` 或 `h200-cu130`）。
+2. 实现 §3.1–§3.6。
+3. 写 §5.1 + §5.2 单元测试。
+4. 在 8×H100 / H200 上跑 §5.3 smoke，对比 evict 频次和 TTFT p99。
+5. 若 §6.2 风险 1 成立，进 SGLang `radix_cache.py` 看是否需要给 streaming-session req 加 `is_session_active=True` flag 阻止"丢弃 decode 尾"。
+
+---
+
+**核心句**：把 session 当 lifecycle 边界（保留），但**不要**让它做 eviction 边界（移交给 radix LRU）。这次 refactor 同时解决"reseed 太频繁"和"P 端 radix 不可外部喂入"两个 blocker。