Retract、decode 内存压力与请求回退#

这章解决什么问题#

前面的调度与内存章节已经解释了 batch 怎样形成、KV cache 怎样复用、allocator 怎样管理页和不变量怎样被检查，但还差一个真实系统里非常关键的问题：当 decode 阶段真的顶到内存压力时，系统不是立刻崩溃，而是怎样把部分请求回退、释放空间、继续把剩下的请求跑下去。

这一章专门讲 retract。如果不补这一层，调度与内存章节仍然会偏“平稳运行时”的解释，而缺少“资源不够时系统怎样自救”的部分。

为什么 retract 比 eviction 更值得单独讲#

eviction 讲的是缓存怎么被回收；retract 讲的是已经在 decode path 上的请求怎样被部分撤回。两者都和内存压力有关，但语义不同：

  • eviction 主要作用在 cache 结构上。
  • retract 直接改变 running batch 里的请求集合和请求状态。

从 schedule_batch.py 和 scheduler.py 看，retract 不只是一个内部 helper，而是 scheduler 在内存不足时主动改变运行队列的正式机制。

一张图：decode OOM 时，系统怎样从“继续跑”切到“回退一部分请求”#

这张图解决的理解障碍是：很多读者会以为 OOM 只会导致“报错并终止”。实际上，这里存在一个有策略的回退路径。

flowchart LR
    Batch["running decode batch"] --> Check["check_decode_mem()"]
    Check -->|enough| Keep["continue decode"]
    Check -->|not enough| Retract["retract_decode()"]
    Retract --> Release["release_req() + release_kv_cache()"]
    Release --> Queue["re-add request to waiting/retracted queue"]
    Retract --> Abort["abort last request if still cannot fit"]

这张图比纯文字多解释的一点是：retract 不是失败后的补丁，而是 decode 调度本身内建的分支。

check_decode_mem() 在检查什么#

schedule_batch.py 里的 check_decode_mem() 先计算 new_tokens_required_next_decode(...)，再调用 evict_from_tree_cache(self.tree_cache, num_tokens)，最后判断 token_to_kv_pool_allocator.available_size() >= num_tokens。

这说明 decode 内存压力判断不是“静态看当前还剩多少页”，而是先估算下一轮 decode 要新增多少 token，再尝试从 tree cache 里腾挪可用空间，然后才做最终判断。也就是说，它检查的是“下一轮还能不能推进”，不是“当前状态看起来紧不紧张”。

retract_decode() 真正做了什么#

retract_decode() 的逻辑很像一本好系统书会重点解释的“压力路径”：

  1. 先准备 sorted_indices。
  2. 如果不是 speculative decoding，就按 (输出长度, -输入长度) 排序，决定优先回退谁。
  3. 在 check_decode_mem(...) 仍不满足时，不断从 batch 里 pop 掉请求并调用 release_req(...)。
  4. 如果最后只剩一条请求仍然塞不下，就给它设置 FINISH_ABORT(...)，而不是让整个 scheduler 直接崩掉。
  5. 最后 filter_batch(...)，并重新计算 new_estimate_ratio。

这里最值得写进书里的，是第四步：即使回退到最后只剩一条请求，系统也会优雅地 abort，而不是把 scheduler 逼进不可恢复状态。这是典型的工程性取舍，优先保住服务可继续运行。

release_req(...) 为什么不是简单 free 掉内存#

release_req(...) 做的事情至少有四步：

  1. 如果启用了 hisparse_coordinator，先通知它 retract 请求。
  2. decode disaggregation 模式下，可能先做 offload_kv_cache(...)。
  3. release_kv_cache(req, self.tree_cache, is_insert=False) 释放 KV，但不把这次回退立刻当作正常 cache insert。
  4. 立刻 evict_from_tree_cache(...)，把刚变得 evictable 的空间尽快转成可用空间。
  5. req.reset_for_retract()，重置请求侧状态。

这说明 retract 不是普通完成路径。它既要尽快回收资源，又要把请求保留成可重新排队的状态，因此既不能完全像 finished request 那样处理，也不能只 free 内存不重置逻辑状态。

reset_for_retract() 为什么值得单独强调#

reset_for_retract() 会：

  • retraction_count += 1
  • 清空或重置 prefix / routed experts / last node / logprob 临时状态 / chunked 标记等
  • 把 is_retracted = True
  • 把 retracted_stain = True

其中最有价值的一个字段是 retracted_stain。它的含义不是“这条请求已经坏了”，而是“这条请求曾经被回退过，因此后续某些 cached token 统计不能再按第一次那样理解”。这就是优秀技术书应该强调的细节：某些字段存在，不是为了功能，而是为了让后续统计和决策不再误读历史。

scheduler 怎样消费 retract 的结果#

scheduler.py 在 decode OOM 路径里会：

  1. 记录 old_available_tokens 和 old_ratio。
  2. 调 batch.retract_decode(self.server_args)。
  3. 记录 num_retracted_reqs，并更新 metrics。
  4. 对必须 abort 的请求，通过 send_to_tokenizer.send_output(AbortReq(...)) 主动发回。
  5. 打 warning，说明这次 retract 回收了多少 token、new_token_ratio 如何变化。
  6. 把 retracted 请求重新 _add_request_to_queue(req, is_retracted=True)。

这说明 retract 并不是 batch 内部悄悄发生的操作，而是会显式改变调度统计、请求排队状态和对外错误输出。

为什么 speculative decoding 下 retract 更难做#

源码注释写得很直接：对 spec decoding，filter_batch 目前只能从后面开始回退，因此 retract policy 还不够理想。也就是说，retract 策略本身会受执行模式限制。

这是一个很典型的 tradeoff：

  • 普通 decode 可以更自由地按策略选择回退对象。
  • spec decode 因为 batch 结构和验证路径更复杂，回退空间受限。

这类信息对读者很重要，因为它说明“回退策略不好”未必是粗心实现，有时是数据结构和执行模式共同逼出来的限制。

这一层最常见的表面现象#

1. 请求偶尔被重新排队，尾延迟突然变长#

这通常不是网络问题，而可能是 decode 内存压力触发了 retract。优先看：

  • retraction_count
  • scheduler warning
  • num_retracted_reqs
  • metrics 里的 retracted input/output tokens

2. 某个请求被直接 abort，错误看起来像内部故障#

当最后一条请求也装不下时，系统会设置 FINISH_ABORT(...) 并以 HTTPStatus.INTERNAL_SERVER_ERROR 结束它。这时根因不是 API 层，而是 decode 内存压力。

3. 服务没崩，但吞吐突然掉一截#

这往往意味着 scheduler 正在通过 retract 勉强维持可运行，而不是处在稳定高效状态。

这套设计的收益与代价#

收益：

  • 避免 decode OOM 直接把 scheduler 打崩。
  • 让系统能在资源紧张时保住一部分在跑请求。
  • 把内存压力转化为显式的回退与重排，而不是隐式错误。

代价：

  • 被回退请求的尾延迟会明显拉长。
  • 请求状态与 metrics 语义会变复杂，需要额外字段如 retracted_stain、retraction_count。
  • speculative decoding 下回退策略受限，未必总是最优。

如果你怀疑问题和 retract 有关，先从哪里看#

建议按这个顺序：

  1. 看 scheduler warning 是否出现 Retract requests 或相关日志。
  2. 看 check_decode_mem() 的判断是否持续失败。
  3. 看 retract_decode() 实际回退了哪些请求。
  4. 看 release_req() 是否走了 offload / release / evict / reset 全链路。
  5. 看被回退请求是否重新入队，还是已经被 abort。

小结#

这一章想让你得到的，不只是“系统会 retract”这个结论，而是一条更可迁移的工程判断：

  • 缓存回收和请求回退不是一回事。
  • 当 decode 内存不够时，系统优先尝试回退部分请求，而不是整体失败。
  • retract 改变的不只是内存占用，还改变请求状态、调度统计和尾延迟。

到这里，调度与内存这一部分就真正覆盖了“资源足够时怎样跑”和“资源不够时怎样活下来”这两种现实。