Prhub

#23427 [HiCache] Prevent move_hybrid_indices from polluting radix-tree node host state

原始 PR 作者 cs-cat 合并时间 2026-04-25 14:27 文件变更 1 提交数 3 评论 6 代码增减 +29 / -9

执行摘要

修复 move_hybrid_indices 污染 radix-tree 状态的 bug

PR #22940 修复了 kernel 后端下 host indices 未移至 CUDA 的崩溃,但引入了新回归:move_hybrid_indices 原地修改共享 PoolTransfer 状态,导致 host eviction 时 MambaPoolHost.free() 使用被污染的 CUDA indices,触发设备不匹配崩溃。Issue #23429 报告了此 bug。

建议精读此 PR,特别是了解如何通过创建新对象而非原地修改来避免共享状态污染的惯用法。这是一个典型的最小化修复案例,值得参考。

讨论亮点

主要讨论来自 Issue #23429 和 PR 评论。hzh0425 批准并表示“Great work! make some changes; will move on after the hicache ci restored”。由于 HiCache CI 暂时不可用,hzh0425 本地验证后合并。无其他争议性讨论。

实现拆解

  1. 修改 move_hybrid_indices 签名与行为hybrid_cache_controller.py):该方法现在返回三元组 (host_indices, device_indices, resolved_pool_transfers)。对于 operation.pool_transfers 中的每个 PoolTransfer,不再直接修改其 host_indicesdevice_indices,而是创建新的 PoolTransfer 实例并收集到新列表 resolved_pool_transfers 中。原始 PoolTransfer 对象保持原样,避免污染 radix-tree 节点状态。
  2. 更新调用方 start_writingstart_loading:这两个方法使用新返回的 resolved_pool_transfers 替代原来的 op.pool_transfers 传入后续操作(如 backup_from_device_all_layerload_to_device_per_layer),确保使用正确的移动后索引。
  3. 代码优化:第二个 commit "Optimize Code" 由 hzh0425 提交,对代码进行了整理优化。
  4. 测试覆盖:本次变更未包含独立的测试文件,但 hzh0425 声称已本地验证通过。
文件 模块 状态 重要度
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py 混合缓存 modified 6.86

关键符号

move_hybrid_indices start_writing start_loading

关键源码片段

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py entrypoint

核心修复文件,修改 move_hybrid_indices 方法及其调用者 start_writing 和 start_loading。

# hybrid_cache_controller.py 中 move_hybrid_indices 的修复核心def move_hybrid_indices(
    self, operation: CacheOperation
) -> tuple[torch.Tensor, torch.Tensor, Optional[list[PoolTransfer]]]:
    host_indices, device_indices = self.move_indices(
        operation.host_indices, operation.device_indices
    )
    resolved_pool_transfers = None
    if operation.pool_transfers:
        resolved_pool_transfers = []
        for transfer in operation.pool_transfers:
            # 为每个 PoolTransfer 创建本地变量,避免修改原对象
            transfer_host_indices, transfer_device_indices = self.move_indices(
                transfer.host_indices, transfer.device_indices
            )
            # 创建一个新的 PoolTransfer 实例,不修改原对象,
            # 因为 radix-tree 节点可能仍引用原始 transfer。
            resolved_pool_transfers.append(
                PoolTransfer(
                    name=transfer.name,
                    host_indices=transfer_host_indices,
                    device_indices=transfer_device_indices,
                    keys=transfer.keys,
                    hit_policy=transfer.hit_policy,
                )
            )
    return host_indices, device_indices, resolved_pool_transfers

评论区精华

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

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

风险与影响

低风险:变更仅涉及 hybrid_cache_controller.py 一个文件,逻辑简单——将原地修改改为返回新副本。主要风险是 resolved_pool_transfers 是否正确传递,以及是否覆盖所有使用 op.pool_transfers 的路径(start_writingstart_loading 均已更新)。由于 HiCache CI 不可用,未能在 CI 中验证,但 hzh0425 已本地验证。

影响范围:仅影响使用 HiCache 混合模型(如 Mamba)的用户,这些用户启用了 host eviction。修复前会导致崩溃(设备不匹配),修复后恢复正常。影响程度中等,因为崩溃是阻断性的,但仅出现在特定配置下。

依赖本地验证(CI 不可用) 无新增测试覆盖

关联 Issue

#22940 [HiCache]Fix hybrid model move_indices
#23429 [Bug] [HiCache] crash when host-pool eviction occurs

完整报告

参与讨论