From 57909f731b0fb325f0313946cf358a23fbe0fd00 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Sat, 6 Jun 2026 00:56:13 +0800 Subject: [PATCH] [PD] Fix KV cache corruption on abort by notifying ongoing prefill (#27372) Signed-off-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 38 ++++++++++++ .../srt/disaggregation/mooncake/conn.py | 60 +++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 94aa7d569..b4ed4f0db 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -877,6 +877,7 @@ class CommonKVReceiver(BaseKVReceiver): self.conclude_state: Optional[KVPoll] = None self.require_staging: bool = False self.init_time: Optional[float] = None + self.abort_notified: bool = False self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) @@ -1062,6 +1063,13 @@ class CommonKVReceiver(BaseKVReceiver): f"in KVPoll.WaitingForInput", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + if ( + not self.abort_notified + and hasattr(self, "bootstrap_infos") + and self.bootstrap_infos is not None + ): + self._send_abort_notification() + self.abort_notified = True return KVPoll.Failed def failure_exception(self): @@ -1079,6 +1087,36 @@ class CommonKVReceiver(BaseKVReceiver): ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.conclude_state = KVPoll.Failed + if ( + not self.abort_notified + and hasattr(self, "bootstrap_infos") + and self.bootstrap_infos is not None + ): + self._send_abort_notification() + self.abort_notified = True + + def _send_abort_notification(self): + for bootstrap_info in self.bootstrap_infos: + # Best-effort notification to prefill side that this request was aborted. + try: + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + with lock: + sock.send_multipart( + [ + b"ABORT", + str(self.bootstrap_room).encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + ] + ) + logger.debug( + f"Sent abort notification for room {self.bootstrap_room} " + f"to {bootstrap_info.get('rank_ip', 'unknown')}:{bootstrap_info.get('rank_port', 'unknown')}" + ) + except Exception as e: + logger.debug( + f"Failed to send abort notification for room {self.bootstrap_room}: {e}" + ) class CommonKVBootstrapServer(BaseKVBootstrapServer): diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 3d4406452..ccae5cc41 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1168,6 +1168,21 @@ class MooncakeKVManager(CommonKVManager): MooncakeRequestStage.MOONCAKE_WORKER_SEND.level, ) + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Failed + ): + logger.debug( + f"Skipping chunk for room {kv_chunk.room} because it has already failed or been aborted" + ) + if self.enable_trace: + kv_chunk.trace_ctx.trace_slice_end( + MooncakeRequestStage.MOONCAKE_WORKER_SEND.stage_name, + MooncakeRequestStage.MOONCAKE_WORKER_SEND.level, + thread_finish_flag=True, + ) + continue + if ( self.enable_staging and staging_strategy is None @@ -1386,6 +1401,44 @@ class MooncakeKVManager(CommonKVManager): handle_staging_rsp(waiting_req_bytes, self.transfer_infos) continue + # Decode-side abort notification: mark room as failed and ACK + 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")) + # No need to abort the room if it has already succeeded + 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}, " + f"marked as Failed" + ) + else: + logger.debug( + f"Received abort notification for room {room_to_be_aborted}, " + f"ignoring (already completed or unknown)" + ) + # Send ACK back to decode endpoint + 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"), + ] + ) + logger.debug( + f"Sent ABORT_ACK for room {room_to_be_aborted} to " + f"{decode_ip}:{decode_port}" + ) + except Exception as e: + logger.debug( + f"Failed to send ABORT_ACK for room {room_to_be_aborted}: {e}" + ) + continue mooncake_session_id = waiting_req_bytes[3].decode("ascii") if room == "None": self.decode_kv_args_table[mooncake_session_id] = ( @@ -1457,6 +1510,13 @@ class MooncakeKVManager(CommonKVManager): self._handle_staging_req(msg) continue + # Prefill acknowledges abort notification + if msg[0] == b"ABORT_ACK": + # TODO(shangming): use this info to implement the deferred release mechanism if needed + ack_aborted_room = int(msg[1].decode("ascii")) + logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}") + continue + bootstrap_room, status, prefill_rank = msg status = int(status.decode("ascii")) bootstrap_room = int(bootstrap_room.decode("ascii"))