KVCache休眠恢复机制与显存复用

数码系统

相机 win10

测评 win11

手机智车

华为 Tesla

小米理想

苹果蔚来

游戏软件

LOL 抖音

原神微信

当前位置：首页

AI教程

KVCache休眠恢复机制与显存复用

热心网友时间：2026-07-06

转载

上一篇将 vLLM、SGLang、TRT-LLM 等推理引擎封装到了 rollout backend 的抽象框架中。今天我们来深入探讨：当 rollout engine 与训练 engine 共享同一块 GPU 时，代码中为何频繁出现 sleep()、resume()、release_kv_cache()、resume_kv_cache() 这类看似“运维操作”的调用？

这里先给出核心结论：sleep 和 resume 并非为了让推理服务实现优雅的上下线，它们实际上是训推共卡场景下的显存状态机——简单来说，就是决定 rollout 的权重、KV cache、prefix cache，以及训练参数和优化器状态，哪些留在 GPU 上，哪些在权重同步或训练更新之前需要临时腾出空间。

先看看共卡时显存冲突的具体情况。看下图时需要注意，训练 engine 和 rollout engine 争夺的并非一个笼统的“GPU 内存”，而是多个生命周期截然不同的对象。

训推共卡时的显存争用

这张图对应 rollout 配置中几个关键参数：gpu_memory_utilization 控制推理引擎分配给 KV cache 的显存比例；free_cache_engine 决定生成完成后是否释放 engine cache；max_num_batched_tokens、max_num_seqs、prompt_length/response_length 则限定了生成阶段的 token 容量上限（可参考 verl/trainer/config/rollout/rollout.yaml:37-79 和 verl/workers/config/rollout.py:188-212）。

1. 生成结束后为何要 sleep

在 RayPPOTrainer.fit() 的主循环中，执行 generate_sequences() 之后，紧接着会调用 self.checkpoint_manager.sleep_replicas()（verl/trainer/ppo/ray_trainer.py:1373-1384）。如果 validation 阶段需要使用 colocated reward model，也会先让 rollout model 进入 sleep，计算完 reward 后，再通过 update_weights() 把 rollout 唤醒（verl/trainer/ppo/ray_trainer.py:553-566）。

下面这张图将这个动作放回一轮 step 中重新审视。它要补充的关键点是：sleep 不是生成阶段的收尾操作，而是从推理状态切换回训练状态的入口。

rollout 显存状态生命周期

CheckpointEngineManager.sleep_replicas() 的注释说得非常直接：sleep all rollout replicas，释放 weight 和 KV cache 在 device 上的内存（verl/checkpoint_engine/base.py:409-412）。在 naive 同进程路径中，ActorRolloutRefWorker.update_weights() 会在 free_cache_engine 开启时先 resume(tags=["weights"])，更新 rollout weights，再 resume(tags=["kv_cache"])，将推理状态恢复至下一轮 rollout 可用状态（verl/workers/engine_workers.py:663-740）。

由此可见，sleep/resume 的设计意图根本不是让 server “暂停一下”，而是让同一批 GPU 在生成、reward/logprob、actor update、weight sync 等阶段之间交替承载不同的内存对象。

2. 同一份合同在不同 backend 中的含义并不相同

第 14 篇提到 backend abstraction 是一个薄抽象，sleep/resume 就是最典型的例证。下面这张图将 vLLM、SGLang、TRT-LLM 的语义放在一起对比。看图时需特别留意 tags 和 level 参数：它们决定了释放的是 KV cache、weights，还是两者都释放。

不同 backend 的 sleep/resume 语义

vLLM adapter 的 resume(tags) 会转化为 server 的 wake_up(tags=...)，release() 会转化为 sleep(level=...)；更新权重后还会清除 prefix cache 并设置 global steps（verl/workers/rollout/vllm_rollout/vllm_rollout.py:150-192）。vLLM async server 在 colocated 模式下会执行 wake_up(tags=["kv_cache", "weights"]) 并重置 prefix cache；hybrid sleep 会根据 LoRA 或 NPU 情况选择不同的 sleep level（verl/workers/rollout/vllm_rollout/vllm_async_server.py:583-606，verl/workers/rollout/vllm_rollout/vllm_async_server.py:891-906）。

SGLang adapter 将 memory occupation 明确拆分成 tags。默认 sleep_level=2 表示释放 weights 和 KV cache；LoRA adapter 模式下 sleep_level=1，只释放 KV cache，保留 base weights（verl/workers/rollout/sglang_rollout/sglang_rollout.py:184-188，verl/workers/rollout/sglang_rollout/sglang_rollout.py:258-285）。SGLang async server 也会根据 hybrid/colocated/standalone 模式，选择释放 ["kv_cache", "weights"] 或只释放 ["kv_cache"]（verl/workers/rollout/sglang_rollout/async_sglang_server.py:430-494）。

TRT-LLM 路径同样实现了 release_kv_cache() 和 resume_kv_cache()；hybrid/colocated sleep 会释放完整 tags，standalone 则跳过 wake/sleep（verl/workers/rollout/trtllm_rollout/trtllm_async_server.py:403-440）。其 adapter 在恢复 KV cache 前还会执行分布式 barrier，避免非 leader rank 的 actor offload 尚未完成就恢复 KV，导致 OOM（verl/workers/rollout/trtllm_rollout/trtllm_rollout.py:395-421）。

3. 非 naive 权重同步仅挪走 KV cache

naive 路径适用于同进程直接同步；但当 backend 走 NCCL/NIXL/HCCL 等 checkpoint engine 时，权重同步策略要精细得多。下面这张图说明了非 naive 路径的执行顺序：先中止并保存未完成请求，再释放 KV cache，保留 weight buffer，完成权重传输，最后恢复 KV cache 和 generation。

checkpoint engine 中的 KV cache 同步路径

CheckpointEngineManager.update_weights() 对这个路径描述得很清楚：非 naive 情况下先 abort_replicas()，再 release_kv_cache_replicas()；注释特别说明它不同于 sleep_replicas()，只释放 KV cache，保留 model weights，目的是让传输直接写入已有的 weight buffers。之后构建 process group、trainer/rollout 一起 update weights、finalize，最后 resume_kv_cache_replicas() 和 resume_generation_replicas()（verl/checkpoint_engine/base.py:447-492）。

这意味着在 weight sync 过程中，KV cache 就是一块“临时挪开的工作区”。如果直接释放 weights，可能需要重建更大的推理状态；如果完全不释放 KV cache，又可能不给权重同步留足显存。这里的工程取舍并非单纯快慢问题，而是显存峰值、同步成本和请求连续性之间的平衡。

4. 读懂 sleep/resume 代码时需避免两个误区

第一，不要把 free_cache_engine=True 简单理解成“更省显存”。它确实能释放推理侧的 cache，但每次生成和训练之间都要付出状态切换的成本。对于短 response、高频 step 的训练场景，这个成本可能非常明显；反过来，对于长 response 或共卡压力大的场景，不释放又可能直接导致 OOM。

第二，不要以为所有 backend 都能以同等粒度释放 KV cache。vLLM、SGLang、TRT-LLM 虽然都暴露了类似的生命周期接口，但 tags、sleep level、prefix cache reset、standalone 模式和 LoRA adapter 路径的实现各不相同。源码中甚至能看到 vLLM async server 的 release_kv_cache() 仍带有 TODO，说明“只释放 KV、保留 weights”这一能力，并非在所有 backend 中都能完全等价实现（verl/workers/rollout/vllm_rollout/vllm_async_server.py:612-623）。

小结：显存复用是 rollout 抽象的一部分

第 13 篇说明 rollout 是训练中的推理服务；第 14 篇说明不同 backend 被抽象成 rollout contract；本文则补充了这层 contract 的显存含义：

generate 需要 weights + KV cache  
→ sleep 释放 rollout 侧显存  
→ actor/reward/logprob 使用训练侧显存  
→ update_weights 恢复 weights  
→ resume_kv_cache 恢复下一轮生成能力

读到这里，rollout 已经不再只是“生成 response 的服务”，而是一个能在训练和推理状态之间灵活切换的 GPU 内存管理边界。下一篇将继续进入 agent loop：当 rollout 从单轮补全变为多轮工具交互时，系统为何必须坚持 token-in/token-out，而不是仅保存最终文本。

本文源码索引

verl/trainer/config/rollout/rollout.yaml:37-79：KV cache、free_cache_engine、batch token 和序列容量配置。
verl/workers/config/rollout.py:188-212：rollout dataclass 中长度、KV 显存比例、free_cache_engine 和容量字段。
verl/trainer/ppo/ray_trainer.py:553-566：validation 中 rollout sleep 后为 colocated reward 腾显存。
verl/trainer/ppo/ray_trainer.py:1373-1384：训练主循环生成后调用 sleep_replicas()。
verl/checkpoint_engine/base.py:409-445：sleep/wake、abort、release/resume KV cache 的 manager 方法。
verl/checkpoint_engine/base.py:447-492：非 naive checkpoint engine 更新权重时释放 KV、同步权重、恢复 KV 和 generation。
verl/workers/engine_workers.py:663-740：naive 权重同步路径中的 resume(weights) 和 resume(kv_cache)。
verl/workers/rollout/vllm_rollout/vllm_rollout.py:150-192：vLLM adapter 的 resume/release/update weights。
verl/workers/rollout/vllm_rollout/vllm_async_server.py:583-623：vLLM server wake/sleep/cache 控制。
verl/workers/rollout/vllm_rollout/vllm_async_server.py:891-906：vLLM hybrid sleep level 选择。
verl/workers/rollout/sglang_rollout/sglang_rollout.py:184-188：SGLang sleep_level 的语义。
verl/workers/rollout/sglang_rollout/sglang_rollout.py:258-285：SGLang adapter 的 resume/release tags。
verl/workers/rollout/sglang_rollout/async_sglang_server.py:430-494：SGLang async server 的 wake/sleep/KV cache 控制。
verl/workers/rollout/trtllm_rollout/trtllm_async_server.py:403-440：TRT-LLM async server 的 cache 和 sleep/wake 控制。
verl/workers/rollout/trtllm_rollout/trtllm_rollout.py:395-421：TRT-LLM adapter 恢复 KV 前的 barrier 和 tags 映射。