15-KVCache-sleep-resume与显存复用

上一篇把 vLLM、SGLang、TRT-LLM 放进了 rollout backend 抽象。本文继续往下看：当 rollout engine 和训练 engine 共用 GPU 时，为什么代码里到处都是 sleep()、resume()、release_kv_cache()、resume_kv_cache()这些看似“运维”的动作？

本文的核心判断是：sleep/resume 不是为了让推理服务优雅上下线，而是训推共卡时的显存状态机。它决定 rollout weights、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 memory（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"])并 reset 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。随后 build process group、trainer/rollout 一起 update weights、finalize，最后 resume_kv_cache_replicas()和 resume_generation_replicas()（verl/checkpoint_engine/base.py:447-492）。

这意味着 KV cache 在 weight sync 里是一块“临时挪开的工作区”。如果直接释放 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 映射。