Prhub

#27372 [PD] Fix KV cache corruption on abort by notifying ongoing prefill

原始 PR 作者 ShangmingCai 合并时间 2026-06-06 00:56 文件变更 2 提交数 10 评论 8 代码增减 +98 / -0

执行摘要

修复 PD 解耦中止时 KV 缓存损坏

Fix KV cache corruption on abort by notifying ongoing prefill. This should address almost all race conditions caused by long ITL aborts. (PR body)

该 PR 值得精读,尤其是设计决策:轻量通知 vs 延迟释放。了解 PD 解耦系统中 abort 处理的权衡对相关开发者有帮助。但需注意代码中未处理的消息解析风险。

讨论亮点

gemini-code-assist[bot] 提出了四点关注:

  • 属性名错误(decode.py):使用 decode_initiated_abort 但实际属性为 abort_initiated,导致延迟释放机制完全失效。作者在后续提交中回退了延迟释放逻辑,此问题不再存在。
  • ABORT 消息解析缺少验证:直接解析可能导致背景线程因畸形消息崩溃,建议添加 try-except。最终代码未采纳,风险仍存在。
  • ABORT_ACK 消息解析类似问题:同样缺乏验证,存在潜在崩溃风险。
  • ABORT_ACK 导致内存泄漏:如果 ACK 延迟到达,_abort_acked_rooms 集合持续增长。该问题随延迟释放回退而不适用,但未来实现需注意。

作者未在最终版本中回应这些评论,其中属性名和泄漏问题已因回退自动解决,而消息解析鲁棒性未处理。

实现拆解

  1. 添加 abort 通知标志和发送函数:在 CommonKVReceiver__init__ 中增加 abort_notified 标志,新增 _send_abort_notification 方法,向所有预填充节点发送 ABORT 消息。(文件:common/conn.py
  2. 在 abort 和超时路径中触发通知:在 CommonKVReceiver.abort()_check_waiting_timeout() 中,在标记房间失败后检查 abort_notified,若未发送则调用 _send_abort_notification。(文件:common/conn.py
  3. 在 transfer_worker 中跳过已失败房间:在 MooncakeKVManager.transfer_worker 中,处理 chunk 前检查房间状态,若已失败则跳过该 chunk 避免后续 RDMA 传输。(文件:mooncake/conn.py
  4. 在 bootstrap_thread 中处理 ABORT 消息:在 MooncakeKVManager.start_prefill_thread 内部循环中,识别 ABORT 消息,将对应房间标记为 Failed,并回复 ABORT_ACK。(文件:mooncake/conn.py
  5. 在 decode_thread 中处理 ABORT_ACK 消息:在 MooncakeKVManager.start_decode_thread 内部循环中,识别 ABORT_ACK 并记录日志(留作将来实现延迟释放)。(文件:mooncake/conn.py
  6. (回退)原计划包含 decode 侧的延迟 KV 缓存释放,但因性能权衡被回退,由 commit "revert deffer" 移除。
文件 模块 状态 重要度
python/sglang/srt/disaggregation/common/conn.py PD 解耦 modified 7.18
python/sglang/srt/disaggregation/mooncake/conn.py PD 解耦 modified 6.96

关键符号

_send_abort_notification abort _check_waiting_timeout transfer_worker bootstrap_thread decode_thread

关键源码片段

python/sglang/srt/disaggregation/mooncake/conn.py core-logic

在 mooncake 传输线程和主循环中添加处理逻辑,防止继续传输已失败房间并处理 ABORT 消息。

# bootstrap_thread 中的 ABORT 处理片段
if room == "ABORT":
    room_to_be_aborted = int(waiting_req_bytes[1].decode("ascii"))
    decode_ip = waiting_req_bytes[2].decode("ascii")
    decode_port = int(waiting_req_bytes[3].decode("ascii"))
    # 仅当房间未成功时才标记失败
    if (room_to_be_aborted in self.request_status
            and self.check_status(room_to_be_aborted) != KVPoll.Success):
        self.update_status(room_to_be_aborted, KVPoll.Failed)
        logger.debug(
            f"Received abort notification for room {room_to_be_aborted}, "
            "marked as Failed")
    else:
        logger.debug(
            f"Received abort notification for room {room_to_be_aborted}, "
            "ignoring (already completed or unknown)")
    # 回复 ABORT_ACK
    try:
        na = NetworkAddress(decode_ip, decode_port)
        self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart([
            b"ABORT_ACK",
            str(room_to_be_aborted).encode("ascii"),
        ])
    except Exception as e:
        logger.debug(f"Failed to send ABORT_ACK for room {room_to_be_aborted}: {e}")
    continue# decode_thread 中的 ABORT_ACK 处理片段
if msg[0] == b"ABORT_ACK":
    ack_aborted_room = int(msg[1].decode("ascii"))
    logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}")
    continue

评论区精华

decode_initiated_abort 属性名错误 正确性

在 decode.py 中使用了 decode_initiated_abort,但 CommonKVReceiver 中属性名为 abort_initiated,导致延迟释放机制完全失效。

结论:作者在后续提交中回退了 decode 侧的延迟释放逻辑,因此该问题不再存在。 · resolved (reverted)

ABORT 消息解析缺少验证 正确性

直接解析 ABORT 消息可能因畸形消息导致背景线程崩溃,建议添加 try-except 和长度验证。

结论:最终代码未添加验证,风险仍存在。 · unaddressed

ABORT_ACK 消息解析缺少验证 正确性

类似地,ABORT_ACK 消息直接解析可能导致背景线程崩溃。

结论:最终代码未添加验证,风险仍存在。 · unaddressed

ABORT_ACK 可能导致内存泄漏 other

延迟释放机制中,如果 ABORT_ACK 迟到,则 _abort_acked_rooms 集合持续增长。

结论:由于延迟释放被回退,该问题目前不适用。但 future work 中需要注意。 · resolved (reverted)

风险与影响

  • 消息解析鲁棒性:在 mooncake/conn.py 的 background thread 中解析 ABORTABORT_ACK 时未做异常处理,畸形消息可能导致线程崩溃,影响整个服务。
  • 部分场景未覆盖:当正在进行的预填充已经调用传输同步(如 mooncake timeout 30s)时,本方案无法阻止,尽管概率极低。
  • 回归风险:改动集中在 abort 路径和传输中断路径,不影响正常流程;但新增的网络协议可能因配置或网络问题导致意外行为。
  • 性能影响:增加一次 ABORT+ACK 消息交互,仅发生在 abort 时,无额外开销。
  • 用户:系统稳定性提升,减少 KV 缓存损坏导致的错误和重启。
  • 系统:轻微增加控制面消息量,无显著性能影响。
  • 团队:简化了 abort 处理方案,避免了复杂的延迟释放机制,但遗留了部分边界情况。
消息解析鲁棒性不足 部分场景未覆盖 回退部分设计

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论