执行摘要

• 一句话：PD拆分decode端支持radix缓存，减少KV传输
• 推荐动作：建议所有从事PD拆分和KV缓存优化的工程师仔细阅读该PR。关键设计决策包括：（1）lock_ref平衡策略及其失败路径处理；（2）将decode prefix handoff迁入kv_manager state以避免调度线程遍历transfer_infos；（3）页面对齐和单调游标维护以防协议错误。该PR为后续扩展（Mooncake支持、批处理eviction、retraction支持）奠定了良好基础。

功能与动机

在PD拆分中，decode节点每轮都要接收完整的KV前缀，导致带宽浪费和TTFT居高不下。通过在decode侧启用radix缓存，共享前缀可被本地缓存，decode只请求新增KV页面，显著降低传输延迟和内存占用。

实现拆解

1. decode调度器前缀匹配与锁定：在pop_preallocated中调用新增的_match_prefix_and_lock，通过match_prefix_for_req匹配decode侧radix树并锁定匹配节点，防止驱逐。获取prefix_len后用于预分配预算计算，并传给_pre_alloc初始化extend输入。

2. prefill端增量传输：修改pop_bootstrapped，通过KV管理器获取decode_prefix_len，计算待发送页数（总页数-前缀页数）。在send_kv_chunk中保持发送游标单调递增，当游标超过分块边界时发送空块。add_transfer_request在kv_indices为空时跳过RDMA传输，但辅助数据仍然发送。TransferInfo新增decode_prefix_len字段，修正is_dummy方法以区分全命中与CP非权威rank。

3. 通用工具函数：在common.py新增kv_to_page_indices、kv_to_page_num、page_align_floor，将原本在disaggregation/utils.py中的函数通用化。新增maybe_cache_unfinished_req包装cache_unfinished_req，支持skip_radix_cache_insert标记。调整release_kv_cache在cache finished时使用该标记。

4. CLI与兼容性：在server_args.py新增--disaggregation-decode-enable-radix-cache布尔标志。decode模式下启用该标志会禁止disable_radix_cache，当前仅允许与nixl后端组合，并对speculative decoding、Mamba/SSM/SWA模型添加断言阻止。

5. 测试覆盖：新增单元测试test_decode_radix_lock_ref.py验证四种transfer场景下lock_ref平衡；新增分布式集成测试test_disaggregation_decode_radix_cache.py在8-GPU H20 CI上多轮缓存命中测试；更新CI脚本安装NIXL。

关键文件：

• python/sglang/srt/disaggregation/decode.py（模块 解码调度；类别 source；类型 core-logic；符号 _match_prefix_and_lock, _pre_alloc, _required_alloc_tokens）: 核心变更：添加_prefix_match_and_lock、修改pop_preallocated和_pre_alloc以集成decode侧radix缓存。
• python/sglang/srt/mem_cache/common.py（模块 缓存工具；类别 source；类型 dependency-wiring；符号 kv_to_page_indices, kv_to_page_num, page_align_floor, maybe_cache_unfinished_req）: 新增通用工具函数（kv_to_page_indices、page_align_floor等）并调整release_kv_cache以支持skip_radix_cache_insert标记。
• python/sglang/srt/disaggregation/nixl/conn.py（模块 传输协议；类别 source；类型 core-logic；符号 pop_decode_prefix_len, should_send_kv_chunk, TransferInfo）: 传输协议扩展：TransferInfo添加decode_prefix_len，add_transfer_request支持空KV传输跳过，修正is_dummy方法。
• python/sglang/srt/server_args.py（模块 配置管理；类别 source；类型 core-logic；符号 _handle_pd_disaggregation）: CLI参数和兼容性验证：新增--disaggregation-decode-enable-radix-cache，禁止与不支持的backends和模型特性组合。
• python/sglang/srt/managers/schedule_policy.py（模块 调度策略；类别 source；类型 refactor；符号 match_prefix_for_req）: 提取match_prefix_for_req函数，供decode调度器和普通调度器共用。
• test/registered/unit/mem_cache/test_decode_radix_lock_ref.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestDecodeLockRefScenarios, _make_cache_with_pools, MockReq）: 单元测试覆盖四种transfer锁平衡场景，验证inc/dec_lock_ref无泄漏。

关键符号：_match_prefix_and_lock, _pre_alloc, _required_alloc_tokens, match_prefix_for_req, kv_to_page_indices, kv_to_page_num, page_align_floor, maybe_cache_unfinished_req, TransferInfo.is_dummy, TransferInfo.from_zmq, pop_decode_prefix_len, should_send_kv_chunk

关键源码片段

python/sglang/srt/mem_cache/common.py

新增通用工具函数（kv_to_page_indices、page_align_floor等）并调整release_kv_cache以支持skip_radix_cache_insert标记。

def kv_to_page_indices(kv_indices: np.ndarray, page_size: int):
    # 页面除最后一块外保证是完整的
    if page_size == 1:
        return kv_indices
    return kv_indices[::page_size] // page_size
​
​
def kv_to_page_num(num_kv_indices: int, page_size: int):
    return (num_kv_indices + page_size - 1) // page_size
​
​
def page_align_floor(length: int, page_size: int) -> int:
    return (length // page_size) * page_size
​
​
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
    # 当标记为跳过 radix 缓存插入时（例如 decode radix cache 场景），直接返回
    if getattr(req, "skip_radix_cache_insert", False):
        return
    tree_cache.cache_unfinished_req(req, **kwargs)

python/sglang/srt/disaggregation/nixl/conn.py

传输协议扩展：TransferInfo添加decode_prefix_len，add_transfer_request支持空KV传输跳过，修正is_dummy方法。

@dataclasses.dataclass
class TransferInfo:
    """传输信息，包含从 decode 发送到 prefill 的索引和参数。"""
    room: int
    endpoint: str
    dst_port: int
    agent_name: str
    dst_kv_indices: npt.NDArray[np.int32]
    dst_aux_index: int
    required_dst_info_num: int
    dst_state_indices: List[int]
    decode_prefix_len: Optional[int] = None # 新增：decode 侧已缓存的前缀长度
​
    def is_dummy(self):
        # CP 非权威 rank 的传输为 dummy；
        # 但当 kv_indices 为空且 decode_prefix_len > 0（全命中）时不是 dummy
        if self.dst_kv_indices.size == 0 and self.decode_prefix_len:
            return False
        return self.dst_kv_indices.size == 0
​
    @classmethod
    def from_zmq(cls, msg: List[bytes]):
        # 解析 state_indices
        dst_state_indices = []
        if len(msg) > 7 and msg[7] != b"":
            dst_state_indices = list(np.frombuffer(msg[7], dtype=np.int32))
        return cls(
            room=int(msg[0].decode("ascii")),
            endpoint=msg[1].decode("ascii"),
            dst_port=int(msg[2].decode("ascii")),
            agent_name=msg[3].decode("ascii"),
            dst_kv_indices=np.frombuffer(msg[4], dtype=np.int32),
            dst_aux_index=int(msg[5].decode("ascii")),
            required_dst_info_num=int(msg[6].decode("ascii")),
            dst_state_indices=dst_state_indices,
            decode_prefix_len=(
                int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None
            ),
        )

评论区精华

• Transfer失败路径lock_ref平衡（@yudian0504）：总结四种场景的lock_ref变化，要求验证所有路径。作者通过单元测试覆盖成功和失败路径，确认无泄漏。（已解决）
• 非NIXL后端兼容性（@ShangmingCai）：询问Mooncake等其他后端是否支持。作者暂时仅允许NIXL并添加ERROR，后续PR添加支持。（已解决）
• Retraction与radix缓存交互（@cctry）：指出retraction需要正确处理前缀匹配，目前未实现。作者添加TODO，计划后续PR添加。（待处理）

• Warmup垃圾KV处理（@cctry）：warmup请求使用伪造bootstrap地址，产生的KV应该被丢弃。作者在warmup结束后flush radix缓存解决。（已解决）

• Transfer失败路径lock_ref平衡 (correctness): 作者通过新增单元测试test_decode_radix_lock_ref.py覆盖四种场景，确认inc/dec_lock_ref平衡且无泄漏。

• 非NIXL后端兼容性 (design): 当前仅允许NIXL后端；Mooncake支持将在后续PR中添加。
• Retraction与radix缓存交互 (correctness): 作者确认未实现，并添加TODO注释，计划后续PR处理。
• Warmup垃圾KV处理 (correctness): 作者通过warmup结束后执行flush_cache清除所有garbage KV（commit a6b6e8b）。

风险与影响

• 风险：
  • lock_ref泄漏风险：如果transfer失败后lock_ref未正确递减，将导致缓存行永久锁定，最终耗尽内存。已通过单元测试覆盖四种场景，但实际部署中仍可能遇到未预见的失败路径。
  • 页对齐不一致：decode和prefill对前缀长度进行页对齐时若算法不同，可能导致传输内容错位。代码在pop_preallocated中使用page_align_floor对齐，并在kv_to_page_indices中处理。
  • 非NIXL后端：当前仅支持NIXL，若其他后端被启用，将因decode_prefix_len参数缺失而崩溃。已在server_args.py中禁止其他后端，但后续添加Mooncake时需额外工作。
  • DP attention命中率：在DP模式下，如果同一对话被路由到不同rank，缓存命中率将大幅下降。需要外部路由器支持prefix-aware路由。
  • eviction未批处理：当前每个请求独立执行eviction，可能在高负载下成为性能瓶颈，计划后续批处理优化。
• 影响：
  • 用户影响：新增CLI选项，用户可在PD拆分decode节点上启用radix缓存，获得显著的TTFT降低和吞吐提升（实测TTFT P50降低8.1倍，吞吐提升1.32倍），但ITL可能因batch变大而退化。
  • 系统影响：decode节点内存利用率提高（更多请求可同时运行），但增加的batch size可能增加解码延迟。对prefill节点影响有限，仅减少发送的KV页面数。
  • 团队影响：核心团队需为Mooncake等后端适配同样功能；调度团队需考虑prefix-aware路由支持；测试团队需维护新增的分布式测试。
  • 风险标记：核心路径变更, lock_ref泄漏风险, 非NIXL后端未支持, eviction未批处理

关联脉络

• 暂无明显关联 PR