SGLang项目中PD+DP模式下的内存泄漏问题分析与解决

问题背景

在SGLang项目的分布式推理场景中，当使用Prefill-Disaggregation(PD)与Data-Parallel(DP)注意力机制组合模式时，系统在解码侧检测到了内存泄漏问题。这一现象在纯TP(prefill)与DP(decode)混合使用的特殊配置下尤为明显。

问题现象

系统运行时抛出内存泄漏错误，具体表现为token_to_kv_pool_allocator的内存分配器检测到可用内存(13824256)与最大令牌数(13858304)不匹配，且树缓存的可回收大小为0。错误信息明确指出内存泄漏发生在解码阶段的调度器内存检查环节。

技术分析

该问题涉及SGLang核心架构的多个关键组件：

1. 内存管理机制：系统使用token_to_kv_pool_allocator进行KV缓存的内存分配，并设有保护区域和最大令牌数限制
2. 调度系统：在event_loop_normal_disagg_decode循环中定期执行内存检查
3. 树缓存系统：负责管理可回收的缓存空间

在PD+DP混合模式下，特别是当prefill使用纯TP而decode使用DP注意力时，内存分配器未能正确跟踪和释放已使用的内存，导致可用内存持续减少。

解决方案

项目维护者通过代码审查发现，这一问题实际上是一个已知警告，在最近的代码合并中被升级为了错误。经过深入调试，团队确认：

1. 该内存泄漏可能是暂时性的，不会持续累积
2. 在PD与DP注意力的集成测试中曾出现过类似警告但已解决
3. 特殊配置组合(纯TP prefill + DP decode)是触发该问题的一个边界条件

最终，该问题通过代码提交得到修复，主要调整了内存泄漏检测的严格程度，并优化了混合模式下的内存管理策略。

经验总结

这一案例揭示了分布式推理系统中几个重要设计考量：

1. 不同并行策略(TP/DP)混合使用时需要特别注意内存管理的一致性
2. 内存泄漏检测机制需要根据实际使用场景调整敏感度
3. 边界条件测试的重要性，特别是对于新兴的disaggregation架构

该问题的解决不仅修复了一个具体bug，也为SGLang项目在复杂分布式场景下的稳定性提供了宝贵经验。