执行摘要

本 PR 通过提取公共基类方法（_start_heartbeat_checker_thread、_handle_node_failure、_prepare_send、_check_bootstrap_timeout、_check_waiting_timeout）和共享数据结构（TransferKVChunk、AuxDataCodec），将 Mooncake、Mori、Nixl 三个 PD 后端中大量重复的逻辑统一到 common/conn.py 和 common/utils.py。各后端删除本地副本并继承基类。重构后代码减少 402 行，新增 283 行，CI 通过，功能无退化。

功能与动机

PR Body 明确指出："The Mooncake, Mori, and Nixl disaggregation backends share significant duplicated logic for heartbeat checking, node failure handling, KV transfer preparation, and timeout management. This duplication increases maintenance burden and makes it easy for bug fixes in one backend to be missed in others." 本次重构正是为解决此问题，将重复逻辑集中到公共基类，让各后端只保留差异化代码。

实现拆解

1. 公共数据结构迁移：在 common/utils.py 新增 TransferKVChunk（传输工作单元）和 AuxDataCodec（辅助数据序列化）。原 Mooncake 和 Nixl 各有本地定义，现统一。
2. 基类方法提取：
   • CommonKVManager._start_heartbeat_checker_thread：循环检测所有注册 prefill 节点的 /health 端点，失败累计超限时调用 _handle_node_failure。成功时调用钩子 _on_heartbeat_success 供子类重写。
   • CommonKVManager._handle_node_failure：清理连接池和预填信息表，将受影响房间标记为 KVPoll.Failed。
   • CommonKVSender._prepare_send：处理发送索引的原子递增、CP rank 过滤，返回过滤后的索引和是否最后一块等指示。
   • CommonKVSender._check_bootstrap_timeout：基于 init_time 检测超时。
   • CommonKVReceiver._check_waiting_timeout：检测接收等待超时。
3. 后端适配：
   • Mooncake：删除本地 TransferKVChunk、AuxDataCodec、heartbeat_checker，改用公共导入；重写 _on_heartbeat_success 以清理成功的 tracker。
   • Nixl：删除本地 TransferKVChunk、heartbeat_checker、_handle_node_failure；移除 TransferStatus.is_failure 字段（已有 request_status 表达故障）。
   • Mori：删除本地 AuxDataCodec 并导入公共类；将参数 is_last 统一为 is_last_chunk。
4. 配置与导入简化：各后端移除 filter_kv_indices_for_cp_rank 的导入（由公共 _prepare_send 调用）。

python/sglang/srt/disaggregation/common/utils.py

公共数据结构定义文件，新增 TransferKVChunk 和 AuxDataCodec，消除各后端的重复定义。

import ctypes
import dataclasses
from typing import List, Optional
​
import numpy as np
import numpy.typing as npt
​
@dataclasses.dataclass
class TransferKVChunk:
    """KV 缓存传输的工作单元：每次传输一个切片。"""
    room: int # 关联的 bootstrap room
    prefill_kv_indices: npt.NDArray[np.int32] # 待发送的 KV 块索引
    index_slice: slice # 当前切片在总索引中的位置
    is_last_chunk: bool # 是否为此传输的最后一块
    prefill_aux_index: Optional[int] # 可选的辅助数据索引
    state_indices: Optional[List] # 可选的状态数据索引
    chunk_id: Optional[int] = None # 分块标识（用于跟踪）
​
class AuxDataCodec:
    """辅助数据的序列化与反序列化工具类。"""
    @staticmethod
    def serialize_data_from_buffer(src_addr, data_length):
        """从源内存地址读取 data_length 字节并序列化为 bytes。"""
        buffer = (ctypes.c_byte * data_length).from_address(src_addr)
        return bytes(buffer)
​
    @staticmethod
    def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data):
        """将 bytes 数据反序列化到指定 GPU 辅助缓冲区。"""
        dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index]
        item_len = kv_args.aux_item_lens[buffer_index]
        dst_addr = dst_aux_ptr + item_len * aux_index
        buffer = (ctypes.c_byte * len(data)).from_address(dst_addr)
        buffer[:] = data

评论区精华

gemini-code-assist[bot] 提出了三点可供改进的建议：

• 性能：阻塞顺序心跳可能在节点数大时延迟处理，建议异步或线程池。
• 正确性：time.time 受系统时钟影响，应使用 time.monotonic。
• 线程安全：heartbeat_failures 无锁保护。

作者 ShangmingCai 对 Nixl is_failure 字段的删除做出解释：该字段已是死代码，故障通过 request_status 和 _handle_node_failure 表达，无需在 TransferStatus 中重复标记。

风险与影响

• 回归风险：心跳失败清理路径的细微差异可能影响预处理节点重连。Mooncake 原有成功心跳时清理 tracker 的逻辑已通过钩子保留，但其他后端原无此行为，需验证是否一致。
• 性能风险：当前心跳仍为单线程阻塞 HTTP，大规模部署时可能成为瓶颈。
• 可维护性提升：后续增加新后端或调整公共策略（超时、重试）只需改动基类一处，团队协作成本降低。

关联脉络

本 PR 是 sglang PD 模块持续重构的一部分。类似地，PR #26134 也在 attention 后端统一 CUDA graph 的 capture/replay 逻辑。之前 PR #25844 扩展了 KV event 的观测能力，PR #26017 优化了 trtllm_mla 的初始化，这些都共用同一套公共传输层框架。本次重构进一步巩固了公共基础，为后续引入新后端（如更多 RDMA 实现）铺平了道路。