# PR #29258 完整报告

- 仓库：`sgl-project/sglang`
- 标题：Fix fixed-size HiCache capacity under PP
- 合并时间：2026-06-26 12:49
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/29258

---

# 执行摘要

- 一句话：修复 PP 下固定 HiCache 容量不同步的问题
- 推荐动作：该 PR 值得合并，修复了 PP 下 HiCache 的一个潜在正确性问题。设计决策明确：使用 MIN reduce 确保所有 stage 容量不超过最小的 stage，与 device KV 容量同步策略一致。建议阅读 `sync_fixed_hicache_size` 函数的实现，理解其安全退避逻辑。

# 功能与动机

修复 PP 模式下 HiCache L2 固定容量时的调度分歧问题。PR body 指出不同 PP stage 因拥有不同层数导致 `size_per_token` 不同，独立转换 GB 到 host token 容量时会产生差异，进而使 radix/HiCache 状态不同，影响调度一致性。例如 GLM-5.2 FP8 tp=4 pp=4 配置下，PP0/PP1 与 PP2/PP3 的 host token 容量相差约 5.26%（8023168 vs 7621952）。

# 实现拆解

1. **新增 `sync_fixed_hicache_size` 函数 **（`python/sglang/srt/mem_cache/pool_host/base.py`）：接收本地计算的 token 容量和 host_size 参数。当 `host_size > 0` 且分布式环境已初始化时，获取 PP group 的 `cpu_group`，通过 `torch.distributed.all_reduce` 执行 `ReduceOp.MIN` 操作，取所有 PP stage 的最小值作为统一的 host token 容量。非 PP 或单 stage 时直接返回原值。
2. **修改 `HostKVCache.__init__`**（`base.py`）：在固定大小分支中，将 `self.size` 的计算包装为 `sync_fixed_hicache_size()` 调用，替换原来直接计算的方式。
3. **修改 `MambaPoolHost.__init__`**（`python/sglang/srt/mem_cache/memory_pool_host.py`）：在 Mamba pool 的 host 容量计算中同样使用 `sync_fixed_hicache_size()`，保持行为一致。同时更新 import 语句，引入 `sync_fixed_hicache_size`。
4. **仅影响初始化路径**：该同步仅在 host pool 构造时执行一次，不影响推理热路径。

关键文件：
- `python/sglang/srt/mem_cache/pool_host/base.py`（模块 缓存层；类别 source；类型 core-logic；符号 sync_fixed_hicache_size, HostKVCache.__init__）: 核心变更文件：新增 `sync_fixed_hicache_size` 函数，并修改 `HostKVCache.__init__` 使其在固定大小分支调用该同步函数。
- `python/sglang/srt/mem_cache/memory_pool_host.py`（模块 缓存层；类别 source；类型 core-logic；符号 MambaPoolHost.__init__）: 修改了 MambaPoolHost 的初始化逻辑，使其在固定大小分支也调用 `sync_fixed_hicache_size`，保证 Mamba host pool 同样获得正确的同步行为。

关键符号：sync_fixed_hicache_size, HostKVCache.__init__, MambaPoolHost.__init__

## 关键源码片段

### `python/sglang/srt/mem_cache/pool_host/base.py`

核心变更文件：新增 `sync_fixed_hicache_size` 函数，并修改 `HostKVCache.__init__` 使其在固定大小分支调用该同步函数。

```python
# 在 base.py 中新增的同步函数，用于对齐 PP stage 间 host token 容量
# 使用 ReduceOp.MIN 确保所有 stage 使用最小公共容量

def sync_fixed_hicache_size(size: int, host_size: int) -> int:
    """Sync fixed-size HiCache token capacity across PP ranks.

    A fixed --hicache-size is specified in GB, but each PP stage may have a
    different bytes/token because it owns different layers. Use the global
    minimum token capacity within the PP group so all stages expose the same
    host-cache capacity.
    Ratio-based sizing already derives from the synced device pool size.
    """
    # 只有在使用固定大小且分布式可用时才同步
    if host_size <= 0 or not torch.distributed.is_available():
        return size

    if not torch.distributed.is_initialized():
        return size

    try:
        from sglang.srt.distributed.parallel_state import get_pp_group

        pp_group = get_pp_group()
    except AssertionError:
        return size

    # 单 stage 不需要同步
    if pp_group.world_size <= 1:
        return size

    tensor = torch.tensor(size, dtype=torch.int64)
    torch.distributed.all_reduce(
        tensor,
        op=torch.distributed.ReduceOp.MIN,
        group=pp_group.cpu_group,
    )
    synced_size = int(tensor.item())

    if synced_size != size:
        logger.info(
            "Sync fixed-size HiCache host token capacity from %d to %d.",
            size,
            synced_size,
        )
    return synced_size

```

```python
# HostKVCache.__init__ 中调用同步函数的位置（partial snippet）
        self.size_per_token = self.get_size_per_token()
        if host_size > 0:
            self.size = sync_fixed_hicache_size(
                int(host_size * 1e9 // self.size_per_token), host_size
            )
        else:
            self.size = int(device_pool.size * host_to_device_ratio)

```

### `python/sglang/srt/mem_cache/memory_pool_host.py`

修改了 MambaPoolHost 的初始化逻辑，使其在固定大小分支也调用 `sync_fixed_hicache_size`，保证 Mamba host pool 同样获得正确的同步行为。

```python
# MambaPoolHost.__init__ 中调用同步函数（partial snippet）
        self.size_per_token = self.get_size_per_token()

        if host_size > 0:
            self.size = sync_fixed_hicache_size(
                int(host_size * 1e9 // self.size_per_token), host_size
            )
        else:
            self.size = int(device_pool.size * host_to_device_ratio)

```

# 评论区精华

Review 中 ShangmingCai 建议将同步范围从全局 group 缩小到 PP group，仅在 pp_size > 1 时执行，避免不必要的跨 TP rank 同步（TP rank 间 `size_per_token` 自然一致）。作者 ziang663 采纳建议，修改为使用 `get_pp_group().cpu_group`，并添加 `world_size <= 1` 的提前返回。

- 同步范围应该限定在 PP group 还是全局 (design): 作者采纳，改为使用 `pp_group.cpu_group`，并在 `world_size <= 1` 时提前返回。

# 风险与影响

- 风险：风险较低。同步仅在初始化时执行一次，不引入运行时开销。使用 `ReduceOp.MIN` 取最小值可能导致部分 stage 实际分配的 host 内存小于 `--hicache-size` 配置值，但业务上这是正确行为（保持容量一致）。若 `get_pp_group()` 抛出 AssertionError（非 PP 环境），函数安全返回原值，无副作用。
- 影响：影响范围限于使用固定 `--hicache-size` 且 PP > 1 的部署场景。修复后 PP stage 间 host cache token 容量一致，调度决策对齐，避免因容量差异导致的性能偏差或异常。对单 stage、非 PP 或使用比例缩放（ratio-based sizing）的情况无影响。
- 风险标记：初始化路径变更 , 分布式同步操作

# 关联脉络

- PR #29044 Fix KV event publisher bind conflict under PP: 同样处理 PP 模式下的 host 端资源同步问题，同属 PP 兼容性修复系列。
- PR #14194 [feature] implement dcp for deepseek_v2: 涉及 device KV 容量的同步，本 PR 的 host 容量同步起到了对称作用。