SGLang KV Cache Retract 机制详解:普通模式 vs PD 分离模式

一条 input 长度逼近 KV cache 容量的请求,在 decode 阶段会触发一次看起来很费时间的 “retract” 交互。本文讲清楚 retract 到底是什么、为什么会被触发、怎么实现的,以及它在普通模式和 PD(prefill-decode)分离模式下走的是两条完全不同的路。

retract 解决什么问题

decode 阶段每生成一个 token,每个跨越 page 边界的请求都要从 KV cache 池子里申请新 page。当池子见底、下一步 decode 放不下时,系统有两种选择:

1. 直接 OOM 崩掉 scheduler;
2. 主动撤回若干个正在 decode 的请求,释放它们的 KV cache,把腾出来的空间留给剩下的请求继续跑。

retract 就是第 2 种。它是 decode OOM 的兜底机制,用”牺牲少数请求”换”整个 batch 不崩”。

SGLang 作者自己对 retract 的评价很直白,在 PR #708 的描述里写着:

Retract decode is bad, nothing good with it.

retract 是个昂贵操作,所以系统还配套了一个 new_token_ratio 预防机制,尽量让 retract 别发生。下面先讲机制本身,再讲预防。

retract 的触发点

入口在 scheduler.py 的 update_running_batch:

1
2
3
4
5
6
7

# scheduler.py:3065
if (kv_full_retract_flag := not batch.check_decode_mem()) or (
    TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0
):
    retracted_reqs, new_token_ratio, reqs_to_abort = batch.retract_decode(
        self.server_args
    )

触发条件是 check_decode_mem 返回 False。这个函数算的是”下一个 decode step 需要新分配多少 token”,再跟池子剩余空间比:

1
2
3
4
5

# schedule_batch.py:2454
def check_decode_mem(self, selected_indices=None):
    num_tokens = self.new_tokens_required_next_decode(selected_indices)
    evict_from_tree_cache(self.tree_cache, num_tokens)
    return self.token_to_kv_pool_allocator.available_size() >= num_tokens

new_tokens_required_next_decode 算的是跨越 page 边界的请求数乘以 page_size。所以触发的不是”来了条新请求”,而是 running batch 累计的请求再 decode 一步就要爆池子。

这解释了为什么”接近容量长度的请求”容易触发:长 input 本身 prefill 就占了大头,进入 decode 后每步都吃 page,池子很快见底。

普通模式:retract = 丢弃 KV,重新 prefill

普通模式(非 PD)下,retract 的代价主要来自”要重新 prefill”。核心流程在 retract_decode (schedule_batch.py:2471):

1. 挑牺牲者

非 spec 模式下按规则排序,优先撤回已生成输出多、输入短的请求:

1
2
3
4
5
6
7
8
9

# schedule_batch.py:2482
if not server_args.speculative_algorithm:
    sorted_indices.sort(
        key=lambda i: (
            len(self.reqs[i].output_ids),
            -len(self.reqs[i].origin_input_ids),
        ),
        reverse=True,
    )

spec 模式因为 filter_batch 只能从尾部删,只能从尾部撤。

2. 逐个释放,直到放得下

调用 release_req,关键在 is_insert=False:

1
2
3
4
5
6
7
8
9

# schedule_batch.py:1591
def release_req(*, req, ...):
    if server_args.disaggregation_mode == "decode":
        req.offload_kv_cache(...)   # 只有 PD decode 模式才走这里
    # 普通模式直接释放,不插回 radix tree
    release_kv_cache(req, tree_cache, is_insert=False)
    num_tokens = remaing_req_count * envs.SGLANG_RETRACT_DECODE_STEPS.get()
    evict_from_tree_cache(tree_cache, num_tokens)
    req.reset_for_retract()

is_insert=False 是普通模式最关键的一行:被撤请求的 KV cache 直接释放,不存进 radix tree。因为系统要立即复用这块内存,存进 tree 反而占地方。代价是:这个请求之后重新 prefill 时基本命中不了 cache,近乎从头算。

reset_for_retract (schedule_batch.py:1443) 把 prefix_indices、last_node、kv_allocated_len、extend_range 等 KV 相关状态全清空,标记 is_retracted=True。

3. 扔回队列等重新 prefill

1
2
3

# scheduler.py:3125
for req in retracted_reqs:
    self._add_request_to_queue(req, is_retracted=True)

普通模式下走的是 _add_request_to_queue 的 NULL 分支,进 waiting_queue,等有空闲显存再重新 prefill。

4. 最后一个还放不下 → abort

如果撤到只剩 1 个请求还放不下,直接 abort 它:

1
2
3
4
5
6
7
8

# schedule_batch.py:2508
if len(sorted_indices) <= 1 and not self.check_decode_mem(...):
    last_req.to_finish = FINISH_ABORT(
        "Out of memory even after retracting all other requests in the decode batch. Aborting the last request.",
        status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
    )
    reqs_to_abort.append(last_req)
    self.release_req(last_idx, 0, server_args)

注意是 FINISH_ABORT,不是放回队列。scheduler 会把这个错误返回给客户端,请求从系统里彻底消失。

PD 分离模式:retract = offload 到 CPU RAM,本地恢复

PD 模式下,retract 走的是完全不同的路径。区别从 release_req 里那行 if disaggregation_mode == "decode" 就开始了:

1
2
3

# schedule_batch.py:1604
if server_args.disaggregation_mode == "decode":
    req.offload_kv_cache(req_to_token_pool, token_to_kv_pool_allocator)

D 节点 retract 时会先把 KV cache 拷一份到 CPU 内存(RAM,不是磁盘),再释放显存:

1
2
3
4
5
6
7
8

# schedule_batch.py:1488
def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
    token_indices = req_to_token_pool.req_to_token[
        self.req_pool_idx, : self.seqlen - 1
    ]
    self.kv_cache_cpu = token_to_kv_pool_allocator.get_cpu_copy(
        token_indices, mamba_indices=self.mamba_pool_idx
    )

这是 PD 模式专门加的优化:把”昂贵的重新 prefill”降级成”GPU↔︎CPU 内存拷贝”。

D 节点本地自恢复,P 节点完全不参与

被撤的请求在 PD 下不进全局 waiting_queue,也不通知 P 节点。它进的是 D 节点本地的 retracted_queue:

1
2
3

# scheduler.py:2324
elif self.disaggregation_mode == DisaggregationMode.DECODE:
    self.disagg_decode_prealloc_queue.add(req, is_retracted=is_retracted)

add 里 is_retracted 分支直接塞进 retracted_queue:

1
2
3
4

# disaggregation/decode.py:485
if is_retracted:
    req.retraction_mb_id = None
    self.retracted_queue.append(req)

等显存有空了,resume_retracted_reqs 把它救回来,全程在 D 节点本地完成:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# disaggregation/decode.py:595
def resume_retracted_reqs(self, rids_to_check=None):
    ...
    for i, req in enumerate(self.retracted_queue):
        if self.req_to_token_pool.available_size() <= 0:
            break
        full_required, swa_required = self._prealloc_required_tokens(req)
        if full_required > full_allocatable_tokens:
            break
        resumed_reqs.append(req)
        indices_to_remove.add(i)
        req.is_retracted = False
        self._pre_alloc(req)
        ...
        # load from cpu, release the cpu copy
        req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator)

load_kv_cache 把之前 offload 到 CPU 的 KV 拷回显存。整个恢复过程 P 节点完全不参与,不需要跨节点重传 KV,这才是 PD 模式那套 retracted_queue 机制存在的意义。

两种模式对比

一句话总结:普通模式 retract 是”丢了重算”,PD 模式 retract 是”暂存到 CPU 内存,本地拷回来”。PD 模式因为 D 节点手里已经有完整 KV(从 P 迁过来的),所以能玩这套;普通模式没有这个前提,只能丢了重 prefill。

为什么 retract 这么慢

普通模式下慢的原因:

• 要重新 prefill:被撤请求的 KV 已释放,得把整条序列重算一遍。prefill 远比单步 decode 贵,长序列尤其痛。
• 重算基本命中不了 cache:撤的时候 is_insert=False 没存进 radix tree,是近乎从头算。PR #440 用 prev_output_ids 保留了已生成 token,重 prefill 时 origin 部分能命中,output 部分还是要重算。
• 连锁撤回:一次可能撤多个请求,每个都要重 prefill。

PD 模式下,retract 本身没那么慢(offload + load 都是拷贝),但前提是 CPU 内存够放得下被撤请求的 KV。如果 D 节点 GPU 显存紧张到频繁 retract,CPU RAM 也可能吃紧,这时 offload 本身也会有压力。

new_token_ratio 预防机制:尽量别触发 retract

retract 太贵,所以 PR #708 引入了提前预防:调度时不按 max_new_tokens 全额预占 KV,而是按 max_new_tokens × new_token_ratio 保守预占。

1
2
3
4
5
6

# schedule_policy.py:872
def add_req_state(r, insert_sort=False):
    new_token_ratio = (
        1.0 if r.sampling_params.ignore_eos else self.new_token_ratio
    )
    tokens_left = r.sampling_params.max_new_tokens * new_token_ratio - len(r.output_ids)

ratio 是动态的,由 NewTokenRatioTracker 管理:

1
2
3
4

# new_token_ratio_tracker.py
init = SGLANG_INIT_NEW_TOKEN_RATIO * schedule_conservativeness   # 默认 0.7
min_ratio = init * SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR              # 默认 ×0.14 ≈ 0.098
decay = (init - min_ratio) / SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS   # 600 步衰减到 min

服务器空闲时 ratio 高(0.7,保守,少接受请求);忙时每步 decay,逐步降到 ~0.098 接受更多请求。一旦真的发生 retract,estimate_new_token_ratio_after_retract 会把 ratio 重设为:

1
2

# new_token_ratio_tracker.py
new_estimate_ratio = (total_decoded_tokens + SGLANG_RETRACT_DECODE_STEPS * len(reqs)) / (total_max_new_tokens + 1)

也就是给剩下的请求留 20 步(SGLANG_RETRACT_DECODE_STEPS=20)的呼吸空间,避免立刻再次 retract。

这个预防机制配合准入 gate(schedule_policy.py:862 的 paged_input > min(cur_rem_tokens, rem_total_tokens) 就返回 NO_TOKEN 拒绝),能在请求被接受阶段就预估它未来要吃多少空间,提前拒绝会撑爆池子的请求。

会不会死循环

一个自然的担心:最后那个请求被 abort 后放回队列,下次推理不又会撞同样的问题?

不会,有三层防护:

第一层:abort 是真终止,不是重试。 最后那个请求被 FINISH_ABORT 后直接返回错误给客户端,从系统里消失,不放回队列。少了这个压力源,下一步通常就放得下了。

第二层:new_token_ratio 主动收紧。 每次 retract 都会调高调度保守度,系统变得更保守,不再接受会撑爆池子的大请求,从源头减少再次触发。

第三层:准入 gate 拦截。 一个超过池子容量的请求根本进不来 running batch,会一直在 waiting_queue 等到池子有空。

但有个极端场景确实”处理不了”:单条请求就超过 KV cache 总容量。这种请求永远过不了准入 gate 卡在 waiting_queue,或者因长度估算不足进了 batch 被 retract 到最后 abort。客户端若重试同一条,又卡回去。这不是 bug 是物理限制,单卡装不下的请求本就无法 decode。解法是增大 KV 容量(更长 context 的卡 / 更多 TP / offload),或 PD 下用更小的 chunked 策略。

控制参数

ServerArgs(最常用)

它乘进 new_token_ratio 生效。注释原话:“Use a larger value if you see requests being retracted frequently.”

环境变量(进阶调优)

版本历史

retract 是 SGLang 早期就有的基础能力,不是 0.5.x 新加的。三个关键 PR:

0.5.x 上围绕 retract 的新 PR(#20724、#23441、#24001、#24675、#26276、#29471、#30899 等)都是给 DSV4 / PD disaggregation / hicache 这些新场景打的补丁,机制本身没变。

调优建议

PD decode 场景下 retract 频繁且拖慢吞吐时:

1. 第一选择:调大 --schedule-conservativeness(1.2~1.5),让调度更早拒绝会撑爆的请求。
2. 其次:调高 SGLANG_INIT_NEW_TOKEN_RATIO,提高初始保守度。
3. 根本解法:给 decode worker 的 KV 容量留足余量,或控制 prefill→decode 迁移时的请求长度,别让单批累计逼近容量上限。

PD 模式下还要注意 CPU RAM 容量:retracted_queue 里 offload 的 KV 都堆在 D 节点内存里,频繁 retract 可能导致内存压力,这时 resume_retracted_reqs 也会变慢。

延伸阅读

PD 模式下真正值得深挖的是 disaggregation/decode.py 里 retracted_queue / resume_retracted_reqs / _swa_aware_allocatable_token_budgets 这套,以及它和 disagg_decode_prealloc_queue 的关系。SWA(sliding window attention)场景下还有 _swa_retractable_len 和 retractable_swa_tokens 的额外预算计算,逻辑比本文讲的更复杂。