# PR #23209 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Refactor] Move radix-cache utils onto RadixKey as methods
- 合并时间：2026-04-21 14:11
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/23209

---

# 执行摘要

- 一句话：将基数树缓存的工具函数重构为 RadixKey 类方法，提升代码封装性和一致性。
- 推荐动作：建议仔细阅读 `radix_cache.py` 中新增的 `RadixKey` 方法实现，特别是 `match` 方法中的 bigram 处理逻辑。这展示了如何将复杂的分支逻辑封装到类方法中，值得学习其设计权衡。对于涉及缓存系统的开发者，此 PR 是理解代码库演进的重要参考。

# 功能与动机

PR body 中明确说明这是一个 ' 纯粹的重构 '，目的是代码整洁。将总是以 RadixKey 为第一个参数的函数转化为类方法，遵循面向对象设计原则，减少代码重复并提高封装性。

# 实现拆解

1. **在 RadixKey 类中添加方法**：在 `python/sglang/srt/mem_cache/radix_cache.py` 中，为 `RadixKey` 类新增 `match`、`child_key`、`hash_page` 和 `_check_compatible` 方法，同时删除模块级的 `_key_match_page_size1`、`_key_match_paged`、`get_child_key` 和 `_check_extra_key` 函数。
2. **更新缓存子类初始化**：在 `swa_radix_cache.py`、`mamba_radix_cache.py`、`unified_radix_cache.py` 等文件中，移除对旧函数的导入，并删除 `__init__` 中设置 `self.key_match_fn` 和 `self.get_child_key_fn` 的逻辑。
3. **修改调用点**：将所有使用 `self.key_match_fn(a, b)` 的地方改为 `a.match(b, page_size=self.page_size)`，将 `self.get_child_key_fn(key)` 改为 `key.child_key(self.page_size)`。
4. **更新测试和文档**：修改 `test_unified_radix_cache_unittest.py` 等测试文件以适配新 API，并更新 `unified_cache_components/README.md` 中的文档引用。
5. **影响**：简化了缓存系统的配置，减少了部分函数的依赖，为后续拆分 `PlainKey`/`BigramKey` 子类铺平道路。

关键文件：
- `python/sglang/srt/mem_cache/radix_cache.py`（模块 缓存核心；类别 source；类型 core-logic；符号 _check_compatible, match, child_key, hash_page）: 定义了 RadixKey 类，本次重构的核心，将自由函数移动为类方法。
- `python/sglang/srt/mem_cache/swa_radix_cache.py`（模块 SWA 缓存；类别 source；类型 dependency-wiring）: SWA 缓存实现，展示了如何移除旧函数依赖并更新为 RadixKey 方法调用。
- `python/sglang/srt/mem_cache/unified_radix_cache.py`（模块 统一缓存；类别 source；类型 dependency-wiring）: 统一缓存实现，同步更新了匹配和子键生成逻辑以使用 RadixKey 方法。

关键符号：RadixKey.match, RadixKey.child_key, RadixKey.hash_page, RadixKey._check_compatible

## 关键源码片段

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

定义了 RadixKey 类，本次重构的核心，将自由函数移动为类方法。

```python
def _check_compatible(self, other: "RadixKey") -> None:
    # 检查 extra_key 是否匹配，确保 RadixKey 操作在相同上下文中进行
    if self.extra_key != other.extra_key:
        raise ValueError(
            f"RadixKey operations require matching extra_key, but got "
            f"{self.extra_key=} != {other.extra_key=}"
        )

def match(self, other: "RadixKey", page_size: int = 1) -> int:
    """
    计算与other共享的逻辑单元前缀长度，结果按page_size向下取整。
    处理bigram模式（用于EAGLE推测解码）和平常模式，适应不同页面大小。
    """
    self._check_compatible(other)
    t0, t1 = self.token_ids, other.token_ids

    if self.is_bigram:
        # 在 bigram 模式下，比较原始 token；L 个匹配 token 意味着 L-1 个匹配 bigram
        i = 0
        for a, b in zip(t0, t1):
            if a != b:
                break
            i += 1
        matched = max(0, min(i - 1, len(self), len(other)))
        # 根据 page_size 对齐匹配长度
        return (matched // page_size) * page_size if page_size > 1 else matched

    if page_size == 1:
        # 页面大小为 1 时，直接比较 token
        i = 0
        for a, b in zip(t0, t1):
            if a != b:
                break
            i += 1
        return i

    # 页面大小大于 1 时，按块比较
    min_len = min(len(self), len(other))
    i = 0
    while i < min_len:
        if t0[i : i + page_size] != t1[i : i + page_size]:
            break
        i += page_size
    return i

def child_key(self, page_size: int = 1):
    """
    生成前page_size个逻辑单元的哈希键，用于在基数树中查找子节点。
    在bigram模式下，键为元组对；extra_key用于命名空间隔离。
    """
    t = self.token_ids
    if self.is_bigram:
        if page_size == 1:
            plain = (t[0], t[1])
        else:
            plain = tuple((t[j], t[j + 1]) for j in range(page_size))
    else:
        plain = t[0] if page_size == 1 else tuple(t[:page_size])
    return plain if self.extra_key is None else (self.extra_key, plain)

def hash_page(self, start: int, end: int, prior_hash: Optional[str] = None) -> str:
    """
    计算逻辑单元[start, end)的SHA256哈希值，用于缓存一致性验证。
    在bigram模式下，输入重叠的(t_i, t_{i+1})字节对。
    """
    hasher = hashlib.sha256()
    if prior_hash:
        hasher.update(bytes.fromhex(prior_hash))
    t = self.token_ids
    if self.is_bigram:
        for j in range(start, end):
            hasher.update(t[j].to_bytes(4, byteorder="little", signed=False))
            hasher.update(t[j + 1].to_bytes(4, byteorder="little", signed=False))
    else:
        for j in range(start, end):
            hasher.update(t[j].to_bytes(4, byteorder="little", signed=False))
    return hasher.hexdigest()

```

# 评论区精华

由于没有 review 评论，主要讨论点在 PR body 中：'Non-goals - `is_bigram` branches inside each method are left as-is; follow-up PR will split into `PlainKey`/`BigramKey` subclasses.' 这表明团队决定暂时保留 bigram 分支，留待后续重构处理。

- Bigram 分支处理策略 (design): 决定在当前 PR 中保留 bigram 分支，后续通过子类拆分来进一步重构。

# 风险与影响

- 风险：主要风险在于重构可能引入细微的逻辑错误，尤其是在匹配和哈希计算的核心路径上。由于变更涉及多个文件和符号，需要确保所有调用点都正确更新了参数传递（特别是 `page_size`）。此外，删除旧函数可能影响依赖这些函数的其他未更新代码，但本 PR 已全面覆盖相关模块。
- 影响：对终端用户无直接影响，API 保持兼容。系统性能应保持不变，但代码结构更清晰，便于未来维护和扩展。开发团队需要熟悉新的 `RadixKey` 方法调用方式，这可能影响后续开发工作流。
- 风险标记：核心路径变更 , 跨模块影响 , 缺少 review 讨论

# 关联脉络

- PR #23107 [Refactor] Replace `page_align_keys` helper with `RadixKey.page_aligned` method: 基础重构，将页面对齐函数移入 RadixKey 类，本 PR 延续此模式，将更多工具函数转化为方法。