执行摘要
本PR修复了SGLang中HiCache模块的内存泄漏问题,通过克隆host indices避免长期引用旧存储。变更仅涉及两处代码修改,但显著提升了系统在启用分层缓存时的稳定性和资源管理能力。
功能与动机
动机:根据Issue #19332,当启用L3 KV缓存(HiCache)时,SGLang在TCP环境下显示内存使用持续增加,甚至导致OOM。根本原因是HiCache在存储alloc()视图时,保留了旧host-pool backing storage的引用,导致内存泄漏。PR body中提供了复现步骤,显示内存从77g增长到102g。
实现拆解
核心修改点:
- 文件
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py:
- 在
write_backup函数中,将node.host_value = host_indices改为node.host_value = host_indices.clone()。
- 在
mamba_backup_commit函数中,同样添加.clone()。
- 文件
python/sglang/srt/mem_cache/hiradix_cache.py:
- 在
write_backup函数中,将node.host_value = host_indices改为node.host_value = host_indices.clone()。
关键逻辑:通过克隆host indices,确保缓存节点持有独立副本,避免共享引用导致的长期内存占用。
评论区精华
review讨论:
- gemini-code-assist[bot]评论:"克隆host indices以防止共享引用副作用。"
- xiezhq-hermann批准变更。
Issue反馈:dcosmos确认修复有效,解决了#19332中的问题。
风险与影响
风险分析:
- 性能开销:
.clone()可能引入额外内存分配和拷贝,轻微影响性能,但权衡修复内存泄漏是必要的。
- 回归风险:变更范围小,回归风险低,但建议部署后监控内存使用。
影响评估:
- 用户影响:修复了启用HiCache时的内存泄漏,提升系统稳定性,减少OOM发生。
- 系统影响:改善资源管理,适用于长时间运行的高并发工作负载。
关联脉络
相关PR:
- PR #21884:移除HiRadixCache中的TTL硬钉功能,同样涉及缓存管理优化。
- PR #21764:修复HiCache缓存命中统计,属于同一功能线的改进。
演进趋势:近期HiCache相关PR频繁,显示团队在持续优化缓存性能和可靠性,本PR是解决内存泄漏的关键一环。
参与讨论