Prhub

#36203 Cleanup duplicate mamba backup helper

原始 PR 作者 ispobock 合并时间 2026-08-25 00:52 文件变更 1 提交数 1 评论 2 代码增减 +0 / -7

执行摘要

删除 schedule_batch.py 中重复的 mamba backup 辅助函数定义

PR body 中的描述较简略,但从变更内容推测,作者发现 _mamba_pool_needing_backup 方法在文件中被重复定义,属于代码清理类改动,目的是消除冗余,提高代码可维护性。

该 PR 是低风险的代码清理,值得快速合入。无需精读,但可关注调度器中 mamba 相关的逻辑,了解备份机制。

讨论亮点

该 PR 没有任何 review 评论和讨论线程,属于简单清理变更,无争议点。

实现拆解

  1. 定位重复定义:在 python/sglang/srt/managers/schedule_batch.py 中,_mamba_pool_needing_backup 方法在约第 1736 行被重复定义,与第 1729 行的定义完全相同。
  2. 删除冗余定义:删除了第二个重复的方法定义(约 7 行),保留了第一个定义,确保逻辑不受影响。
  3. 验证关联调用:该方法被 offload_kv_cache 方法调用,用于判断是否需要备份 mamba pool,删除重复定义后调用关系保持不变。
文件 模块 状态 重要度
python/sglang/srt/managers/schedule_batch.py 调度器 modified 5.2

关键符号

_mamba_pool_needing_backup

关键源码片段

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

删除了重复定义的 _mamba_pool_needing_backup 方法,是本次变更的唯一文件。

# python/sglang/srt/managers/schedule_batch.py
# 该片段展示保留的唯一定义,以及其调用方 offload_kv_cachedef _mamba_pool_needing_backup(self, req_to_token_pool, allocator):
    # 若 CPU 副本已包含 mamba 状态,则无需额外备份
    if allocator.get_kvcache().cpu_copy_carries_mamba:
        return None
    # 仅对混合 token 池(HybridReqToTokenPool)进行 mamba 池备份
    if not isinstance(req_to_token_pool, HybridReqToTokenPool):
        return None
    return req_to_token_pool.mamba_pooldef offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
    token_indices = req_to_token_pool.req_to_token[
        self.req_pool_idx, : self.seqlen - 1
    ]
    # 调用上述辅助函数,确定是否需要备份 mamba 池
    mamba_pool = self._mamba_pool_needing_backup(
        req_to_token_pool, token_to_kv_pool_allocator
    )
    self.retraction_backup = RetractionBackup(
        cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy(
            token_indices, mamba_indices=self.mamba_pool_idx
        ),
        mamba_cpu=(
            mamba_pool.get_cpu_copy(self.mamba_pool_idx.unsqueeze(0))
            if mamba_pool is not None and self.mamba_pool_idx is not None
            else None
        ),
    )

评论区精华

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

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

风险与影响

该改动仅删除代码,没有逻辑变化,风险极低。但需注意:若未来有人基于被删除的重复定义修改代码,可能会产生冲突,不过这属于常规维护风险。此外,该文件属于调度器核心路径,任何改动都需谨慎,但本次删除不影响行为。

对用户和系统无影响,仅是代码清理。对团队而言,减少了代码冗余,提高了可读性。该改动只影响一个文件,影响范围很小。

低风险 代码清理 核心路径文件

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论