Runtime checker、SessionAwareCache 与内存不变量#

这章解决什么问题#

前面四章已经把 scheduler、cache reuse、allocator 和 batch 策略讲得比较深了，但还缺一层特别像“维护者手册”的内容：系统如何确认这些复杂状态没有悄悄漂移？如果 streaming session 介入之后，prefix cache 又怎样在不改写整个调度流水线的前提下接住会话态 KV？

这一章的任务就是把两个看似分散、实际高度相关的机制放到一起：

  1. scheduler_runtime_checker_mixin.py 如何用 pool invariant 和 watchdog 给复杂调度路径兜底。
  2. SessionAwareCache 如何把 streaming session 的 KV 生命周期包装进原有 prefix cache。

把它们放在一起的原因很简单：两者都在回答“状态复杂之后，系统怎样仍然知道自己没乱”。

为什么这一层比表面看起来更重要#

调度与内存章节到了后半程，读者最容易出现的错觉是：只要主流程能跑，缓存和 allocator 迟早会对。真实系统恰恰相反。很多严重问题都不是“完全跑不起来”，而是：

  • 内存漏了一点点，跑久了才炸。
  • session 状态没清干净，下一轮 streaming 复用了错误 KV。
  • watchdog 发现 forward 卡住，但你不知道该先怀疑 scheduler、cache 还是 detokenizer。

因此，好书需要把这些兜底机制讲出来。否则读者只会得到“快乐路径”，而不是一套能带去现场的判断框架。

一张图：pool invariant 为什么是调度层的最后防线#

这张图要解决的理解障碍是：available、evictable、protected、session_held、uncached 这些字段分散在不同对象里，读者不容易意识到它们最后会被压成一个守恒关系。

flowchart LR
    Avail["available"] --> Eq["pool invariant"]
    Evict["evictable"] --> Eq
    Protect["protected"] --> Eq
    Session["session_held"] --> Eq
    Uncached["uncached"] --> Eq
    Eq --> Total["== total"]
    Eq --> Leak["mismatch => leak / warning / assert"]

相比纯文字，这张图多解释了一层：runtime checker 不是“又写了一堆 debug function”，而是把分散在 scheduler、tree cache、allocator、session wrapper 里的状态重新收敛成少数几个守恒式。

_check_pool_invariant(...) 其实在做什么#

scheduler_runtime_checker_mixin.py 里最核心的函数之一是 _check_pool_invariant(...)：

total_accounted = available + evictable + protected + session_held + uncached
leak = total_accounted != total

这段逻辑短得几乎容易被忽略，但它正是全章的中心。因为调度系统的复杂度并不只是“有很多队列”，而是“同一批 token 可能处在可用、可驱逐、被保护、被 session 占用或尚未真正进缓存等多个状态里”。只要这几个量加起来不守恒，系统就出现了泄漏或状态错账。

_check_full_pool(...)、_check_swa_pool(...)、_check_mamba_pool(...)#

mix-in 不是只做一个统一检查，而是按 pool 类型拆开：

  • _check_full_pool(...) 检查 full pool。
  • _check_swa_pool(...) 检查 SWA 相关 pool。
  • _check_mamba_pool(...) 在 hybrid SSM 场景下额外给出 page-level leak diagnosis。

这说明 runtime checker 不是抽象装饰层，而是和缓存实现变体绑定的。也就是说，系统一旦支持更多 cache 形态，checker 也必须跟着增长，而不是永远沿用一套通用断言。

self_check_during_busy() 与 on_idle() 分别守什么#

这一点特别值得写进书里，因为它能显著提升读者的调试能力。

忙碌时检查#

self_check_during_busy() 会在 last_batch 非空、且某些 speculative 限制满足时，计算 uncached 大小，再调用 _check_full_pool(...)。这里的重点不是“忙的时候也检查”，而是系统承认忙时状态并不完全落稳，因此检查逻辑必须考虑当前 batch 里尚未真正 cache 完的部分。

空闲时检查#

on_idle() 则更像一次完整 housekeeping：

  1. _check_all_pools(...)
  2. _check_req_pool()
  3. _check_tree_cache()
  4. 记录 idle metrics
  5. 发布 KV event
  6. reset token ratio
  7. sleep until next event

这表明 idle 不是“什么都不做”，而是系统做全量 coherence check 的时机。很多长跑稳定性问题只有在这里才会被暴露出来。

create_scheduler_watchdog(...) 把什么信息暴露给你#

create_scheduler_watchdog(...) 返回 WatchdogRaw，它会用：

  • forward_ct 作为计数器
  • scheduler 是否初始化/当前是否有 batch 作为活跃判定
  • _check_all_pools(...) 的信息作为 dump 内容

也就是说，scheduler watchdog 不是只告诉你“卡住了”，而是尝试在卡住时把 batch 与 pool 状态一起吐出来。对工程调试来说，这比单纯打印“线程超时”高价值很多。

SessionAwareCache 为什么是 decorator，而不是改写所有 cache#

session_aware_cache.py 的类定义已经把动机写得很直接：它是一个 BasePrefixCache decorator，专门管理 streaming session KV，避免对调度流水线做侵入式修改。

这是一种非常像好系统设计书会强调的做法：新能力不是先改主流程，而是先找一个稳定边界把它包起来。

match_prefix(...) 的关键判断#

SessionAwareCache.match_prefix(...) 先看请求是否是 streaming session：

  • 如果不是，直接 pass through 到 inner.match_prefix(...)。
  • 如果是，就根据 session_id 找 SessionSlot。
  • 如果 slot 存在且 req_pool_idx 有值，就 restore_to_req(req)，并基于已提交的 KV 长度计算 prefix_len 与 device_indices，最后返回 MatchResult(...)。

这段逻辑的价值在于：它把 streaming session 的前缀复用变成“会话槽位恢复”问题，而不是让主 scheduler 到处写 session 分支。

为什么 SessionSlot 不会在 restore 后立刻清空#

源码注释写得很明确：chunked prefill 时，请求可能因为 budget 等原因在这一轮被 scheduler 拒绝，然后下一轮重试。如果 restore 后立刻把 slot 清空，下轮重试就失去幂等恢复能力。

这件事很适合写成书里的“容易混淆点”：

  • 你看到 slot 没被清空，可能会以为这是泄漏。
  • 但对 chunked prefill + streaming session 来说，这是为了保住 retry 语义。

也就是说，某些看起来像“没及时释放”的状态，其实是为了支持幂等重试而被故意保留。判断它是不是 bug，不能只看状态存在与否，还要看它是否仍能被后续请求正确消费。

scheduler 如何把 SessionAwareCache 接进来#

scheduler.py 在初始化 tree cache 时，如果 server_args.enable_streaming_session 为真，就直接：

self.tree_cache = SessionAwareCache(self.tree_cache)

这个装配点很关键。它说明 streaming session 不是另起一套调度器，而是在原有 tree cache 外套一层 wrapper。收益是主 scheduler loop 基本不需要重写；代价是 cache 层承担了更多“会话态解释”的复杂度。

这一层的典型故障与排障入口#

1. 长时间运行后内存慢慢偏掉#

优先看：

  • _check_all_pools(...)
  • _check_req_pool()
  • _check_tree_cache()
  • scheduler watchdog dump

如果是 hybrid SSM 场景，_check_mamba_pool(...) 给出的 leaked page 集合会非常关键。

2. streaming session 下一轮复用了错误上下文#

优先看：

  • SessionAwareCache.match_prefix(...)
  • SessionSlot.save_from_req(...) / restore_to_req(...)
  • session_controller.py
  • 是否处在 chunked prefill 重试路径

3. forward 卡住，但你不确定是模型执行还是状态错账#

先看 scheduler watchdog 的 dump 内容。如果 pool invariant 已经不守恒，先处理状态错账；如果 invariant 仍守恒，再往 execution model 和 attention backend 深入。这个顺序很重要，因为它能避免你在错误层级上浪费时间。

这套设计的收益与代价#

收益：

  • 用 invariant 让复杂状态重新变得可检查。
  • 用 decorator 让 streaming session 复用已有 prefix cache 体系。
  • 让空闲时全量检查、忙时局部检查、watchdog 超时时 dump 信息三者形成维护闭环。

代价：

  • checker 自身必须跟着 cache 变体成长，维护成本不低。
  • session wrapper 让 cache 层语义变重，阅读难度上升。
  • 某些“状态暂不释放”的设计很容易被误判成泄漏，需要结合运行路径理解。

小结#

这一章真正想交给读者的，是一个比“知道有哪些函数”更可迁移的判断框架：

  • 当调度与缓存越来越复杂时，必须把状态重新压回少数不变量。
  • 当新能力要接入现有流水线时，优先找稳定边界做 wrapper，而不是先撕开主循环。
  • 排障时先确认 invariant 是否还成立，再决定要不要深入执行层。

到这里，调度与内存这一部分就不只是在讲“系统如何高效”，也开始讲“系统如何知道自己还没悄悄坏掉”。