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.
15 KiB
Block-level Eviction Refactor — 设计文档
日期:2026-05-12
前置: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 的语义改成:
- Streaming-session decode 输出在 turn finish 时 增量 commit 进 radix tree。
SessionSlot退化为纯 metadata(仅持last_node+ lock_ref 状态),不再独占 KV 区间。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 铺路)。
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 ~12Kkv_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:
@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:
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 维护:
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 指针:
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:
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:
# 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 §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 测试
@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 关键风险
-
inner.cache_finished_req对 streaming-session req 的兼容性:当前 SGLang 标准 radix 假设 req 在 cache_finished_req 时是 "完整 prefill+decode 完成"。streaming-session 的 req 在每个 turn 结束时还会留下"未完成的 conversation",要确保 inner 在插入时不会把 decode-only tokens 当成可丢弃尾巴。需要 auditradix_cache.py:cache_finished_req的实现。 -
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必须先到。 -
chunked-prefill retry:见 I3。SGLang 当前
restore_to_req不清 slot 字段就是为此 retry。refactor 后必须确认 inner radixmatch_prefix在 retry 下也幂等(标准 radix tree 是的,但要写测试明确锁住这个性质)。
7. 与 D→P sync 工作的关系
block-level evict 是 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 的最小动作
- fork 一个
feat/block-level-evict分支(从improve/audit-and-foundations或h200-cu130)。 - 实现 §3.1–§3.6。
- 写 §5.1 + §5.2 单元测试。
- 在 8×H100 / H200 上跑 §5.3 smoke,对比 evict 频次和 TTFT p99。
- 若 §6.2 风险 1 成立,进 SGLang
radix_cache.py看是否需要给 streaming-session req 加is_session_active=Trueflag 阻止"丢弃 decode 尾"。
核心句:把 session 当 lifecycle 边界(保留),但不要让它做 eviction 边界(移交给 radix LRU)。这次 refactor 同时解决"reseed 太频繁"和"P 端 radix 不可外部喂入"两个 blocker。