Prhub

#31321 [Bugfix] Release Mamba cache after PP dynamic chunk profiling

原始 PR 作者 lixuwei2333 合并时间 2026-07-16 02:32 文件变更 1 提交数 1 评论 2 代码增减 +3 / -1

执行摘要

修复 PP 动态分块 profiling 时 Mamba 缓存泄漏

PP 动态分块 profiling 阶段会构造多达 128 个合成 prefill 请求,对于混合 Mamba 模型,每个请求会同时分配 KV 缓存和 Mamba 缓存槽位。但之前的 cleanup 只调用了 free(req) 释放请求池条目和 KV 缓存,未调用 free_mamba_cache,导致 Mamba 缓存槽位持续被占用,最终可能达到 --max-mamba-cache-size 的上限,使动态分块被禁用或 server warmup 失败。PR body 中明确给出了复现命令和错误日志。

推荐合并。这是一个典型的资源泄漏 bugfix,改动量小、逻辑清晰、复现路径明确,且通过了 PP 相关 CI 测试。

讨论亮点

该 PR 的 review 过程非常简洁,未产生实质性技术讨论。ShangmingCai 在 Issue 评论中执行了 /rerun-group pp 触发 PP 相关 CI 重跑,结果通过。ShangmingCai 最终给予了 LGTM 批准。

实现拆解

  1. python/sglang/srt/managers/scheduler_pp_mixin.pyprofile_and_init_predictor 方法中,定位到 profiling 循环结束后的 cleanup 区域(第 717-724 行)。
  2. 在释放 KV 缓存之后、调用 self.req_to_token_pool.free(req) 之前,插入对 req.mamba_pool_idx is not None 的判断:若不为 None,则调用 self.req_to_token_pool.free_mamba_cache(req) 释放该 profiling 请求占用的 Mamba 缓存槽位。
  3. 确保释放顺序正确:先释放 KV 缓存和 Mamba 缓存,再释放请求池条目,避免悬挂指针或重复释放。
  4. 该改动仅影响 profiling 初始化阶段的 cleanup 逻辑,不影响稳态推理路径。
文件 模块 状态 重要度
python/sglang/srt/managers/scheduler_pp_mixin.py 调度器 modified 5.28

关键符号

profile_and_init_predictor free_mamba_cache

关键源码片段

python/sglang/srt/managers/scheduler_pp_mixin.py core-logic

核心改动文件,在 PP 动态分块 profiling cleanup 路径中增加了 Mamba 缓存释放逻辑。

# 以下代码位于 profile_and_init_predictor 方法的 profiling 循环末尾
# 此处在释放 KV 缓存后、释放请求池条目之前,增加了对 Mamba 缓存的释放
# Release KV and Mamba cache
if req.req_pool_idx is not None:
    kv_indices = self.req_to_token_pool.req_to_token[
        req.req_pool_idx, : req.extend_range.end
    ]
    self.token_to_kv_pool_allocator.free(kv_indices)
    # 新增:如果该请求分配了 Mamba 缓存槽位,则释放
    if req.mamba_pool_idx is not None:
        self.req_to_token_pool.free_mamba_cache(req)
    self.req_to_token_pool.free(req)
    req.kv = None

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险极低。改动仅 3 行,位于 profiling 循环 cleanup 区域内,条件判断 req.mamba_pool_idx is not None 保证了仅在混合 Mamba 模型场景下触发,且释放顺序符合资源管理约定。未引入新的依赖或状态变更。主要风险在于若将来 free_mamba_cache 实现有变(如添加引用计数),此处未考虑,但当前实现是直接的槽位释放,无副作用。

影响范围限定于 PP 动态分块 + 混合 Mamba 模型的初始化阶段。修复后,profiling 期间 Mamba 缓存不再泄漏,动态分块能正确完成初始化,server 正常启动。对稳态推理性能无影响。对非 Mamba 模型或非 PP 配置无任何行为变化。

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论