Prhub

#20697 Fix VRAM leak in overlap scheduling with structured output (#20640)

原始 PR 作者 Cishoon 合并时间 2026-03-23 08:07 文件变更 2 提交数 5 评论 5 代码增减 +17 / -0

执行摘要

修复在启用重叠调度和结构化输出时的 VRAM 泄漏问题。

根据Issue #20640,在启用extra_buffer(重叠调度)和结构化输出时,VRAM会缓慢但持续增长,导致生产环境几小时后出现OOM错误和模型重启。问题仅在这两个条件同时满足时触发,影响系统稳定性。

建议工程团队精读此PR,重点关注闭包环境下GPU张量生命周期的管理策略,可作为异步调度中内存优化的参考案例。

讨论亮点

审核中,主要讨论点围绕VRAM泄漏的机制。hnyls2002询问为何VRAM持续增长而非仅延迟一轮释放,Cishoon解释在tp_worker.py中,重叠调度创建的delay_sample_func闭包捕获了张量并保持引用,导致无法及时释放。双方确认修复方案合理,无重大争议。

实现拆解

实现方案涉及两个关键文件:在model_runner.py_preprocess_logits方法中,设置sampling_info.vocab_mask = None以在应用掩码后立即释放GPU张量;在scheduler.pylaunch_batch_sample_if_needed方法中,设置batch_result.delay_sample_func = Nonelogits_output.next_token_logits = None,以断开闭包对forward_batchlogits_output的引用,从而促进内存回收。

文件 模块 状态 重要度
python/sglang/srt/managers/scheduler.py scheduler modified 7.0
python/sglang/srt/model_executor/model_runner.py model_executor modified 7.0

关键符号

launch_batch_sample_if_needed _preprocess_logits

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

VRAM 泄漏机制讨论 正确性

hnyls2002 询问 VRAM 为何持续增长而非延迟释放,Cishoon 解释闭包引用导致张量无法及时回收。

结论:双方确认修复方案有效,解决了引用持有问题。 · 已解决

风险与影响

风险较低,变更仅涉及引用置空,不改变核心业务逻辑。但需注意潜在副作用:如果下游代码错误依赖这些已释放的张量(如vocab_masknext_token_logits),可能引发异常;尽管测试显示内存稳定,但缺少单元测试覆盖具体场景。

直接影响是解决了生产环境中的内存泄漏问题,提高了系统可靠性和性能。影响范围限于使用重叠调度和结构化输出的场景,对普通用例无影响,但有助于提升整体资源利用率。

引用管理变更 缺少测试覆盖

关联 Issue

#20640 [Bug] Qwen3.5+extra_buffer video memory leak during structured JSON response generation.

完整报告

参与讨论