Prhub

#30827 feat: add cache salt support to KV cache events

原始 PR 作者 jthomson04 合并时间 2026-08-13 07:14 文件变更 43 提交数 4 评论 36 代码增减 +754 / -71

执行摘要

为 KV 事件与 radix 缓存新增 cache_salt 命名空间隔离

PR body 明确指出:外部 KV-aware 路由器需要请求级缓存命名空间,且该命名空间必须与 SGLang 调用方自定义的 extra_key 保持区分,因为 Concatenating the two values can collide,而现有 KV 事件哈希与 payload 无法暴露类型化缓存命名空间。cache_salt 因此被设计为与 extra_key 独立的一等字段,用于隔离进程内 radix 缓存并命名空间化外部 KV 事件。

值得精读,尤其关注三点设计决策:(1) wire 兼容策略——用同 tag 独立 struct 追加第 8 槽位而非给基类加可选字段,是 msgspec 生态中可复用的模式;(2) 加盐哈希链的迭代化 + 节点记忆化 + 分裂切分,规避了递归深度与 O(T²) 重算;(3) fail-loudly 设计——直接属性访问与 C++ 后端显式拒绝。对下游路由端,需先实测旧 7 槽位解码器对 8 槽位事件的兼容行为再升级。

讨论亮点

评审中最重要的交锋来自作者自评与 6 条 inline review:

  • getattr 静默降级是隔离特性最糟的失败模式Req 始终定义 cache_salt,应直接 req.cache_saltgetattr(req, "cache_salt", None) 会把“代码路径忘记携带 salt”变成“静默不加盐、跨命名空间共享缓存”,缺失属性应当大声失败。
  • 事件哈希递归重算存在性能与栈溢出双重风险:chunked prefill 每个 chunk 的 store 事件都会把 root→node 整条路径重算一遍(约 O(T²/chunk_size));page_size=1 且共享前缀很深时递归深度可超过 Python 默认 1000 帧限制触发 RecursionError。建议迭代化并在节点上记忆化。
  • store 事件中 parent 链被重复计算_record_store_event 对 parent 的哈希链在同一个事件里算了两遍,记忆化后自然消除。
  • 删除 _compute_extra_key 丢失 eager 校验:list 形式的 cache_salt 会流进声明为 Optional[str]TokenizedGenerateReqInput,在 msgspec IPC decode 深处报晦涩错误而非干净 400,需要在单请求入口补标量校验。
  • C++ 实验后端静默丢盐radix_cache_cpp.py 仍构造 RadixKey(token_ids, req.extra_key),开启 SGLANG_EXPERIMENTAL_CPP_RADIX_TREE=1 时不同 salt 相同 token 的请求会共享缓存;最终选择显式拒绝而非透传。
  • KVEventBatch union 需要注释BlockStoredWithMetadata 刻意不进 union,否则 msgspec 因重复 tag 拒绝定义;已加注释防未来“顺手修复”。

以上意见全部落地;唯一保留的开放项是 L3/远端存储键不按 cache_salt 命名空间的既有限制。

实现拆解

  1. 请求入口与归一化(io_struct.py)GenerateReqInputTokenizedGenerateReqInput 新增 cache_salt 字段;_normalize_single_inputsextra_keycache_salt 做标量字符串校验并把空串归一为 None;新增 _normalize_cache_salt 处理批量展开与并行采样复制;__getitem__ 切片逐项透传;同时收紧 _normalize_extra_key 对列表元素类型与空串的处理,避免两键混淆。

  2. 服务入口与转发(serving_base.py / serving_chat.py / serving_completions.py / encode_receiver.py):删除 _compute_extra_key——此前 chat 路径把 cache_salt 拼进 extra_key,正是碰撞来源;现在 chat/completions/responses 各自把 cache_salt 作为独立字段转发;encode_receiver.create_req 顺带修复了此前完全丢弃 extra_key 的行为(行为变更 bugfix 搭车)。

  3. Radix 键与缓存变体(radix_cache.py / swa_radix_cache.py / mamba_radix_cache.py / unified 系列 / flexkv)RadixKey 增加 cache_salt slot,切片保留;child_key 对加盐键返回 ((extra_key, cache_salt), plain) 结构化二元组,未加盐保持原样;各缓存变体的插入/匹配路径从 req.cache_salt 直接取值(评审后不再用 getattr 静默降级);C++ 实验后端通过 _reject_cache_salt 显式拒绝加盐请求,绝不静默丢盐。

  4. 事件哈希命名空间(mem_cache/utils.py / mem_cache/events.py / unified_tree_core.py):新增 compute_node_event_hash_values:以 SHA256(b"sglang-cache-salt-v1\0" + cache_salt) 为根种子,迭代自顶向下填充路径上缺失的 event_hash_value,节点级记忆化,分裂时复用 split_node_hash_value 切割;_record_store_event / _record_remove_event 对加盐节点改用事件哈希并写入 BlockStoredMetadata

  5. 事件 wire 格式(disaggregation/kv_events.py):新增 BlockStoredMetadataBlockStoredWithMetadata(共享 BlockStored tag,kw_only 追加第 8 槽位);未加盐事件保持 7 槽位旧布局;KVEventBatch 的 tagged union 刻意不引入新类型,现有消费方按基类解码并忽略尾随元数据。

  6. 测试与 CI 配套:新增/扩展 test_io_struct、test_radix_cache_unit、test_mem_cache_utils、test_kv_events(wire 兼容)、test_encode_receiver、test_serving_completions、test_radix_cache_cpp_unit 等,覆盖归一化、碰撞隔离(拼接相等但 salt 不同的反例)、跨节点分裂哈希保持、1100 层路径迭代、msgspec round-trip 与 C++ 后端拒绝;CI 重跑 radix_cache / hicache / disaggregation 组直至通过。

文件 模块 状态 重要度
python/sglang/srt/mem_cache/utils.py 缓存工具 modified 7.29
python/sglang/srt/mem_cache/radix_cache.py 缓存核心 modified 6.94
python/sglang/srt/managers/io_struct.py 请求归一 modified 7.15
python/sglang/srt/disaggregation/kv_events.py 事件协议 modified 6.88
python/sglang/srt/mem_cache/events.py 事件记录 modified 6.71
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py 统一缓存 modified 6.91
python/sglang/srt/mem_cache/radix_cache_cpp.py 缓存后端 modified 6.42
python/sglang/srt/entrypoints/openai/serving_base.py 服务入口 modified 6.45
python/sglang/srt/managers/schedule_policy.py 调度策略 modified 5.93
test/registered/unit/mem_cache/test_radix_cache_unit.py 缓存测试 modified 7.19
test/registered/unit/mem_cache/test_mem_cache_utils.py 缓存测试 modified 6.96
test/registered/unit/disaggregation/test_kv_events.py 事件测试 modified 6.8
test/registered/unit/disaggregation/test_encode_receiver.py 编解码测试 added 6.59

关键符号

compute_node_event_hash_values _normalize_cache_salt RadixKey.__init__ RadixKey.child_key _record_store_event _record_remove_event _reject_cache_salt prefetch_anchor_info _compute_extra_key match_prefix_for_req

关键源码片段

python/sglang/srt/mem_cache/utils.py core-logic

新增 compute_node_event_hash_values:加盐事件哈希的唯一实现点,以 SHA256 根种子派生命名空间哈希链、迭代计算并节点记忆化,是隔离契约能否成立的核心。

# python/sglang/srt/mem_cache/utils.pydef compute_node_event_hash_values(node: Any, page_size: int) -> List[str]:
    """计算并缓存面向外部 KV 事件的命名空间哈希。    内部缓存哈希(compute_node_hash_values)只覆盖 token 序列,
    不同 cache_salt 的请求会得到相同的内部哈希,因此这里用
    cache_salt 派生独立根种子,再沿 radix 树做增量哈希链。
    """
    cache_salt = node.key.cache_salt
    # 未加盐请求直接复用内部哈希,保持旧行为完全不变
    if cache_salt is None:
        return compute_node_hash_values(node, page_size)
​
    # 已算过则直接返回:否则 chunked prefill 每个 chunk 都会
    # 重算已插入前缀,长请求退化为 O(T^2 / chunk_size)
    if node.event_hash_value is not None:
        return node.event_hash_value
​
    # 迭代向上收集尚未填充 event_hash_value 的祖先节点。
    # 递归版本在 page_size=1、共享前缀很深时会超过 1000 帧
    # 限制触发 RecursionError,必须用显式栈
    missing_nodes = []
    current = node
    while (
        current is not None
        and current.key is not None
        and len(current.key) > 0
        and current.event_hash_value is None
    ):
        # 同一条路径上混用不同 cache_salt 会使哈希链失去意义,宁可报错
        if current.key.cache_salt != cache_salt:
            raise ValueError("Radix path contains mismatched cache_salt values")
        missing_nodes.append(current)
        current = current.parent
​
    if (
        current is not None
        and current.key is not None
        and len(current.key) > 0
        and current.key.cache_salt != cache_salt
    ):
        raise ValueError("Radix path contains mismatched cache_salt values")
​
    # 根种子:SHA256("sglang-cache-salt-v1\0" + cache_salt),
    # 保证不同命名空间产生完全独立的哈希链
    if current is not None and current.event_hash_value:
        parent_hash = current.event_hash_value[-1]
    else:
        parent_hash = hashlib.sha256(
            b"sglang-cache-salt-v1\0" + cache_salt.encode("utf-8")
        ).hexdigest()
​
    # 自顶向下逐节点计算并记忆化;节点分裂时直接复用
    # split_node_hash_value 按页切分,与内部 hash_value 同构
    for missing_node in reversed(missing_nodes):
        hash_values = get_hash_str(missing_node.key, parent_hash, page_size=page_size)
        assert isinstance(hash_values, list)
        missing_node.event_hash_value = hash_values
        if hash_values:
            parent_hash = hash_values[-1]
​
    assert node.event_hash_value is not None
    return node.event_hash_value
python/sglang/srt/mem_cache/radix_cache.py core-logic

RadixKey 增加 cache_salt 维度并改写 child_key 结构化索引,所有 Python radix 变体(含 swa/mamba/unified/flexkv)的隔离都依赖这里。

# python/sglang/srt/mem_cache/radix_cache.pyclass RadixKey:
    __slots__ = ("token_ids", "extra_key", "cache_salt", "is_bigram", "limit")
​
    def __init__(
        self,
        token_ids: array[int],
        extra_key: Optional[str] = None,
        is_bigram: bool = False,
        limit: Optional[int] = None,
        cache_salt: Optional[str] = None,
    ):
        self.token_ids = token_ids
        # extra_key:调用方自定义的请求分类键(如 lora_id)
        self.extra_key = extra_key
        # cache_salt:缓存命名空间,必须与 extra_key 独立存放,
        # 否则拼接会产生碰撞;同时用于进程内 radix 树与外部
        # KV 事件的命名空间化,远端 L3 存储键不在此契约内
        self.cache_salt = cache_salt or None
        self.is_bigram = is_bigram
        self.limit = limit
​
    def child_key(self, page_size: int = 1):
        # ...(bigram / 单 token 页的 key 抽取逻辑略)
​
        # 加盐后按结构化二元组 (extra_key, cache_salt) 索引,
        # 未加盐请求保持原有索引行为,向后兼容
        if self.cache_salt is not None:
            return ((self.extra_key, self.cache_salt), plain)
        return plain if self.extra_key is None else (self.extra_key, plain)
python/sglang/srt/disaggregation/kv_events.py data-contract

定义 BlockStoredMetadata / BlockStoredWithMetadata 与 KVEventBatch union 约束,决定外部路由器可见的 wire 契约。

# python/sglang/srt/disaggregation/kv_events.pyclass BlockStoredMetadata(msgspec.Struct, omit_defaults=True, gc=False):
    """附加在加盐 BlockStored 上的类型化请求元数据。"""
​
    cache_salt: str
​
​
class BlockStored(KVCacheEvent):
    # 7 槽位旧布局:未加盐事件必须保持原样,外部消费者按
    # 位置解析 block_hashes / parent_block_hash / token_ids 等
    block_hashes: list[int]
    parent_block_hash: Optional[int] = None
    token_ids: list[int]
    block_size: int
    lora_id: Optional[int] = None
    medium: Optional[str] = None
​
​
class BlockStoredWithMetadata(BlockStored, tag="BlockStored", kw_only=True):
    """仅在存在类型化元数据时使用的 wire 扩展。    独立 struct 而非给 BlockStored 加可选字段,是为了让未加盐
    事件序列化后仍为 7 个元素;若加可选字段,msgspec 会为
    未加盐事件写出尾随 null,破坏旧消费者。
    """
​
    metadata: BlockStoredMetadata
​
​
class KVEventBatch(EventBatch):
    # BlockStoredWithMetadata 刻意不加入这个 tagged union:
    # 现有类型化消费方按共享的 "BlockStored" tag 解码为基类并
    # 跳过尾随元素;若把两个类型都放进 union,msgspec 会因
    # 重复 tag 直接拒绝定义。将来想“修复”时必须先读这条注释。
    events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]

评论区精华

cache_salt 取值应 fail loudly 而非 getattr 静默降级 正确性

作者在 schedule_policy.py 的 diff 上指出:Req 始终定义 cache_salt,直接使用 req.cache_salt 即可;getattr(req, "cache_salt", None) 会把“代码路径漏传 salt”变成“静默不加盐、跨命名空间共享缓存”,对隔离特性是最糟的失败模式,并点名 radix_cache.py、mamba_radix_cache.py、swa_radix_cache.py 等 7 处同类 getattr。

结论:改为直接访问 req.cache_salt,缺失属性时抛 AttributeError;最终所有缓存变体均直接透传 cache_salt。 · 已解决

事件哈希递归重算:性能与 RecursionError 风险 性能

初版 compute_node_event_hash_values 递归重算 root→node 整条路径且无记忆化:chunked prefill 每个 chunk 的 store 事件都会重算已插入前缀,长请求约 O(T²/chunk_size);page_size=1 且共享前缀很深时递归深度可超 1000 帧触发 RecursionError。建议迭代化并在节点上记忆化。

结论:重写为迭代自顶向下 walk + 节点级 event_hash_value 记忆化;分裂时用 split_node_hash_value 同步切分,测试覆盖 1100 层路径与 memoize。 · 已解决

store 事件中 parent 链被重复计算 性能

events.py 的 _record_store_event 先 compute_node_event_hash_values(node) 已隐含算过 parent 链,随后又对 node.parent 调用一次取 parent_block_hash,加盐路径每个 store 事件把 root→parent 走两遍。

结论:随节点记忆化落地后重复计算被消除。 · 已解决

chat 路径删除 _compute_extra_key 后丢失 eager 类型校验 正确性

serving_chat.py 直接用 request.cache_salt 转发后,原来 _compute_extra_key 对非字符串抛 TypeError 的校验消失;ChatCompletionRequest.cache_salt 是 Optional[Union[List[str], str]],列表值会一路流入 TokenizedGenerateReqInput.cache_salt(Optional[str]),最终在 msgspec IPC decode 深处报出晦涩错误而非干净的 400。

结论:io_struct._normalize_single_inputs 对 extra_key / cache_salt 增加 str 类型校验并把空串归一为 None;新增 test_cache_salt_normalization 覆盖。 · 已解决

C++ 实验 radix 后端静默丢弃 cache_salt 正确性

作者自评发现 radix_cache_cpp.py 三处仍以 RadixKey(token_ids, req.extra_key) 构造键;本 PR 之前 salt 是经 _compute_extra_key 拼进 extra_key 的,该后端此前有隔离,改成独立字段后被静默丢弃,SGLANG_EXPERIMENTAL_CPP_RADIX_TREE=1 下不同 salt 相同 token 的请求会共享缓存。

结论:最终选择显式拒绝:radix_cache_cpp.py 新增 _reject_cache_salt,启用实验 C++ 后端时对加盐请求抛错,配套 test_cache_salt_is_rejected_without_loading_cpp_extension。 · 已解决

KVEventBatch union 刻意排除 BlockStoredWithMetadata 需要注释 documentation

现有类型化消费方按共享的 BlockStored tag 解码并跳过尾随元素,若将新类型加入 union,msgspec 会因重复 tag 直接拒绝;作者要求写注释防止未来“顺手修复”。

结论:已添加说明注释。 · 已解决

encode 路径捎带修复 extra_key 丢失 正确性

encode_receiver.py 的 create_req 此前整体丢弃 recv_req.extra_key,导致 keyed 请求与未 keyed 请求共享缓存;本次新增 extra_key 转发是行为变更的 bugfix 搭车,要求写入 PR 描述并配测试。

结论:已补充 extra_key / cache_salt 双转发与新增 test_encode_receiver.py 单测。 · 已解决

L3/ 远端存储未按 cache_salt 命名空间(既有限制) 设计

内部哈希 get_native_hash 只覆盖 token ids,hicache storage / LMCache / FlexKV 远端块键基于未加盐链,加盐请求 miss 设备 / 主机树后仍可能预取其他 salt 写入的远端块;extra_key 同样有此缺口,非本 PR 回归,但既然 PR 正式化 salt 隔离契约,值得在文档中说明。

结论:作为既有限制记录在案,未在本 PR 内修复;PR body 明确 cache-salt-plus-LoRA 及远端语义不在范围内。 · 待处理

风险与影响

  • 兼容性风险(wire 格式):加盐事件第 8 槽位对按 7 槽位解码的旧版外部路由器构成格式变更。测试覆盖了“8 槽位可由类型化消费方 round-trip”与“未加盐保持 7 槽位”,但未覆盖“旧版 7 槽位解码器遇到 8 槽位”的行为,msgspec 对定长数组可能直接报错,升级需与下游路由端协同。
  • 性能风险:加盐请求且开启 KV 事件时,首次访问需沿根路径填充整条 event_hash_value 链(已记忆化且仅一次),长前缀请求首次开销仍存在;内部缓存哈希 get_native_hash 路径不受影响。
  • 回归风险child_key 返回结构变化影响全部 radix 变体;_normalize_extra_key 行为微调(空串转 None、元素类型校验)可能影响依赖旧行为的调用方;_compute_extra_key 删除改变了 chat 路径对非法类型的校验方式(已由 _normalize_single_inputs 补齐)。
  • 功能边界SGLANG_EXPERIMENTAL_CPP_RADIX_TREE=1cache_salt 同时使用时直接报错(显式拒绝而非降级);远端 L3/hicache 存储键未加盐,加盐请求在设备/主机树 miss 后理论上仍可能预取其他命名空间的远端块(既有限制)。
  • 外部路由端:KV-aware 路由器首次能从 BlockStored 事件拿到类型化 cache_salt 元数据,可按命名空间做定向清理与路由,不再依赖危险的手工拼接。
  • 用户侧Engine.generate / OpenAI 兼容接口新增可选 cache_salt 参数,纯增量;未加盐请求完全走旧路径,无行为变化。
  • 系统侧:radix 树隔离粒度提升,hash 链按命名空间独立;但 KV 事件新增第 8 槽位属于对下游消费方的格式演进,需要协同升级。
  • 团队侧:43 个文件的跨模块改动 + 大量单测,为后续 cache-salt-plus-LoRA 路由语义预留了正式契约。
核心缓存路径变更 事件 wire 格式扩展 跨模块参数透传(43 文件) C++ 后端显式拒绝加盐 哈希链首次计算开销 远端 L3 键未加盐 旧消费方 8 槽位兼容未实测

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论