Prhub

#39226 [Bugfix] Fix workspace resize leaking reserved GPU memory

原始 PR 作者 czhu-cohere 合并时间 2026-04-24 04:50 文件变更 1 提交数 7 评论 4 代码增减 +18 / -21

执行摘要

修复 MoE workspace 动态调整时的显存泄漏问题

移除默认 chunking 后,profiling 阶段不再有最坏情况分配,MoE workspace 会在推理过程中动态扩容。当前实现虽先删除旧 buffer 再分配新 buffer,但 CUDA 缓存分配器可能不立即回收,导致 2x 显存驻留。PR body 指出:'sometimes we will have 2x the buffer resident in memory which may cause OOM or cause less space to be reserved for KV cache'。

值得精读。该 PR 解决了实际部署中遇到的显存泄漏问题,展示了在动态形状推理中管理 GPU 显存的典型权衡。建议关注 empty_cache 的性能影响,并评估后续是否在初始化阶段预留最坏情况 workspace 以避免运行时调整。

讨论亮点

gemini-code-assist[bot]: 'Calling torch.cuda.empty_cache() inside a performance-critical path...can cause significant performance degradation due to the global synchronization of the CUDA allocator.' 建议文档化性能瓶颈或用更高效的内存管理策略。
LucasWilkinson (approver): 'Makes sense to me; I think we should add the worst case allocation back during init though if possible.' 倾向于在初始化时分配最坏情况大小以避免运行时调整。
czhu-cohere: 回复说明无 chunking 时最坏情况分配可能过大(如 256k tokens),实际上不可行。

已解决的问题:显存泄漏与峰值。未解决:empty_cache 的性能开销;初始化时分配最坏情况容量的长期方案。

实现拆解

  1. 删除循环批量调整逻辑:原代码循环 for ubatch_id in range(self._num_ubatches) 对所有 ubatch 同时调整 workspace 大小。新代码改为只调整当前请求的 ubatch,其余 ubatch 延迟到下次调用时再调整,避免 DBO 下其他 ubatch 仍持有旧 tensor 视图导致泄漏。
  2. 显式调用 empty_cache:在删除旧 tensor 后、分配新 tensor 前,插入 torch.accelerator.empty_cache() 调用,强制 CUDA 缓存分配器回收刚释放的显存,使新分配能复用它,防止峰值翻倍。
  3. 调整调试日志:日志从报告所有 ubatch 的总内存改为仅报告当前 ubatch 的调整信息,并输出 ubatch_id。

涉及的唯一文件是 vllm/v1/worker/workspace.py,核心方法是 get_simultaneous 内的 workspace 分配逻辑。

文件 模块 状态 重要度
vllm/v1/worker/workspace.py Workspace 管理 modified 6.79

关键源码片段

vllm/v1/worker/workspace.py core-logic

核心变更文件,修改了 workspace 的动态调整逻辑,修复显存泄漏并优化 DBO 场景。

# 仅调整当前请求的 ubatch,其他 ubatch 延迟调整以避免 DBO 下视图泄漏
self._current_workspaces[ubatch_id] = None
del current_workspace# 显式释放已回收的显存段,使缓存分配器能重新使用 GPU 内存进行更大的分配
# 没有这一步,每次调整可能留下未释放的段,导致峰值内存升高
torch.accelerator.empty_cache()# 分配新 workspace
self._current_workspaces[ubatch_id] = torch.empty(
    (required_bytes,), dtype=torch.uint8, device=self._device
)# 更新局部引用,供后续返回
current_workspace = self._current_workspaces[ubatch_id]

评论区精华

empty_cache 性能开销 性能

gemini-code-assist[bot] 指出 empty_cache 导致 CUDA 全局同步,可能引起性能退化,建议文档化或用更高效策略。

结论:作者接受该 trade-off,但未采用替代方案。已在 PR body 中说明 empty_cache 耗时约 5ms(5GB buffer),认为对于大模型 10GB+ 的 buffer 来说,内存节省更重要。 · 已解决

初始化时预留最坏情况 workspace 设计

LucasWilkinson 建议在初始化时分配最坏情况 workspace 以避免运行时调整。czhu-cohere 回应无 chunking 时最坏情况可能过大(如 256k tokens),不可行。

结论:当前方案作为临时修复接受,后续可能重新考虑初始化分配策略。 · unresolved

风险与影响

  1. 性能风险torch.accelerator.empty_cache() 会触发 CUDA 全同步,每次 resize 约耗时 ~5ms(5GB buffer),在高频 resize 场景下可能累积影响。当前仅当 workspace 需要扩容时才触发,频率较低,但仍需监控。
  2. 兼容性风险:使用了 torch.accelerator.empty_cache()(torch 2.0+ 的通用加速器 API),如果项目仍使用较老 PyTorch 版本可能不存在此 API。但 vLLM 已要求较新版本,风险低。
  3. 回归风险:修改了 workspace 调整的循环逻辑,从同步调整所有 ubatch 改为仅调整当前 ubatch。依赖原行为(如预分配所有 ubatch)的代码可能受影响。但 DBO 分支已通过 ubatch_id 区分,且新逻辑通过延迟调整避免泄漏,经过验证。
  4. 测试覆盖:无直接对应的单元测试,依赖集成测试和手动内存剖析验证。

影响范围:使用 MoE 模型且使用 vLLM v1 engine 的用户,尤其是大 EP(专家并行)和 DBO(双批次重叠)场景。
用户影响:降低 OOM 风险,允许更大的 KV cache 预留,可能提升吞吐。代价是 resize 时少量延迟(微秒到毫秒级)。
系统影响:修改了 core workspace 管理逻辑,但仅限于 v1 路径,不影响 v0 或其他后端。
团队影响:验证了动态 workspace 方案的可行性,为后续优化(如初始化时预留)提供基线。

性能同步开销 无直接测试覆盖 动态路径变更

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论