fix(disagg/mooncake): notify prefill on decode-side abort

Port the decode->prefill half of upstream sgl-project/sglang #27372 (the worker
skip-Failed guard half landed in cb1a03f0a). In multi-prefill/multi-decode PD, a
decode-initiated abort frees the decode's KV pages back to the allocator, but the
prefill never learns of it (request_status is per-process), so the prefill keeps
RDMA-writing the remaining chunks into pages the decode may have reallocated to a
different live request -> KV corruption. The already-ported worker guard is a no-op
here because nothing sets the prefill room to Failed.

Now the decode receiver sends a 4-field b"ABORT" notification to every prefill peer
(on abort() and on the poll() WaitingForInput timeout, at most once); the prefill
marks the room Failed (so the worker guard fires) and replies b"ABORT_ACK". The
ABORT branch is handled before the unconditional waiting_req_bytes[3] decode in
bootstrap_thread (a 4-field ABORT would otherwise crash the else branch), and
ABORT_ACK before the 3-tuple unpack in the decode thread.

De-entangled from the staging/tracing infra this branch lacks. Component-tested
(test_mooncake_abort_protocol.py: format, send-once guard, abort wiring, ZMQ
round-trip). Cross-process corruption-window closure to be verified by the
multi-P/D PD harness.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

@@ -1224,6 +1224,34 @@ class MooncakeKVManager(CommonKVManager):
             while True:
                 waiting_req_bytes = self.server_socket.recv_multipart()
                 room = waiting_req_bytes[0].decode("ascii")
+                # Decode-side abort notification: mark the room Failed and ACK, so the
+                # transfer worker (which skips Failed rooms) stops RDMA-writing into KV
+                # pages the decode has freed/reused. Port of upstream sgl-project/sglang
+                # #27372 (decode->prefill half). MUST be handled here, before the
+                # UNCONDITIONAL waiting_req_bytes[3] decode below: a 4-field ABORT
+                # message [ABORT, room, ip, port] would otherwise fall through into the
+                # else branch and crash on waiting_req_bytes[7].
+                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)
+                    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
                 mooncake_session_id = waiting_req_bytes[3].decode("ascii")
                 if room == "None":
                     self.decode_kv_args_table[mooncake_session_id] = (
@@ -1261,6 +1289,12 @@ class MooncakeKVManager(CommonKVManager):
                     self._handle_aux_data(msg)
                     continue
 
+                if msg[0] == b"ABORT_ACK":
+                    # Prefill acknowledged our abort; the room is already Failed locally,
+                    # so nothing more to do. Port of upstream #27372. Must precede the
+                    # 3-tuple unpack below (a 2-field ACK would otherwise raise).
+                    continue
+
                 bootstrap_room, status, prefill_rank = msg
                 status = int(status.decode("ascii"))
                 bootstrap_room = int(bootstrap_room.decode("ascii"))
@@ -1617,6 +1651,7 @@ class MooncakeKVReceiver(CommonKVReceiver):
         self.session_id = mgr.get_session_id()
         self.conclude_state = None
         self.init_time = None
+        self.abort_notified = False
         super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank)
 
         self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
@@ -1746,6 +1781,8 @@ class MooncakeKVReceiver(CommonKVReceiver):
                             self.bootstrap_room,
                             f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput",
                         )
+                        # Tell prefill peers to stop transferring into our pages.
+                        self._send_abort_notification()
                         self.conclude_state = KVPoll.Failed
                         return KVPoll.Failed
 
@@ -1777,6 +1814,35 @@ class MooncakeKVReceiver(CommonKVReceiver):
             )
         raise KVTransferError(self.bootstrap_room, failure_reason)
 
+    def _send_abort_notification(self):
+        # Notify every prefill peer that this room is aborting, so the prefill stops
+        # RDMA-writing into KV pages this decode is about to free/reuse. Without this,
+        # the prefill never learns of a decode-initiated abort (its request_status is
+        # per-process) and the worker's skip-Failed guard never fires. Port of upstream
+        # sgl-project/sglang #27372 (decode->prefill half). Sent at most once.
+        if self.abort_notified:
+            return
+        bootstrap_infos = getattr(self, "bootstrap_infos", None)
+        if not bootstrap_infos:
+            return
+        self.abort_notified = True
+        for bootstrap_info in bootstrap_infos:
+            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"),
+                        ]
+                    )
+            except Exception as e:
+                logger.debug(
+                    f"Failed to send ABORT for room {self.bootstrap_room}: {e}"
+                )
+
     def abort(self):
         self.kv_mgr.record_failure(
             self.bootstrap_room,
@@ -1786,6 +1852,8 @@ class MooncakeKVReceiver(CommonKVReceiver):
         # transfer worker and other pollers observe the abort. Port of upstream
         # sgl-project/sglang #24522.
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
+        # Tell the prefill peers to stop transferring into our (soon-freed) pages.
+        self._send_abort_notification()
         # Explicitly set the status to failure since this request has been aborted
         self.conclude_state = KVPoll.Failed
 

test/registered/unit/disaggregation/test_mooncake_abort_protocol.py

@@ -0,0 +1,138 @@
+"""Component tests for the decode->prefill abort protocol (port of upstream #27372).
+
+These cover the pieces that do NOT need a live PD transfer:
+  - the decode side builds the exact 4-field b"ABORT" message and sends it at most once,
+  - MooncakeKVReceiver.abort() invokes the notification and marks the room Failed,
+  - the ZMQ wire format round-trips so the prefill handler's waiting_req_bytes[1..3]
+    indices are always present (the branch-ordering trap: a 4-field ABORT must not
+    IndexError before the room=="ABORT" check).
+
+The actual cross-process corruption-window closure is verified by the multi-P/D PD
+harness, not here.
+"""
+
+import unittest
+
+from sglang.srt.disaggregation.base.conn import KVPoll
+from sglang.srt.disaggregation.mooncake.conn import MooncakeKVReceiver
+
+
+class _DummyLock:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        return False
+
+
+class _CapturingSock:
+    def __init__(self):
+        self.sent = []
+
+    def send_multipart(self, parts):
+        self.sent.append(parts)
+
+
+class _FakeMgr:
+    def __init__(self):
+        self.local_ip = "10.0.0.1"
+        self.rank_port = 12345
+        self.status_updates = []
+
+    def record_failure(self, room, reason):
+        pass
+
+    def update_status(self, room, status):
+        self.status_updates.append((room, status))
+
+
+def _make_receiver(bootstrap_infos):
+    r = MooncakeKVReceiver.__new__(MooncakeKVReceiver)  # bypass heavy __init__
+    r.abort_notified = False
+    r.bootstrap_room = 777
+    r.bootstrap_infos = bootstrap_infos
+    r.kv_mgr = _FakeMgr()
+    r.conclude_state = None
+    return r
+
+
+class TestSendAbortNotification(unittest.TestCase):
+    def test_message_format_is_four_fields(self):
+        sock = _CapturingSock()
+        r = _make_receiver([{"peer": 1}])
+        r._connect_to_bootstrap_server = lambda info: (sock, _DummyLock())
+        r._send_abort_notification()
+        self.assertEqual(len(sock.sent), 1)
+        parts = sock.sent[0]
+        # [ABORT, room, ip, port] — exactly 4 fields so the prefill handler's
+        # waiting_req_bytes[3] (port) access never IndexErrors.
+        self.assertEqual(len(parts), 4)
+        self.assertEqual(parts[0], b"ABORT")
+        self.assertEqual(parts[1], b"777")
+        self.assertEqual(parts[2], b"10.0.0.1")
+        self.assertEqual(parts[3], b"12345")
+
+    def test_sent_at_most_once(self):
+        sock = _CapturingSock()
+        r = _make_receiver([{"peer": 1}])
+        r._connect_to_bootstrap_server = lambda info: (sock, _DummyLock())
+        r._send_abort_notification()
+        r._send_abort_notification()
+        self.assertEqual(len(sock.sent), 1)
+        self.assertTrue(r.abort_notified)
+
+    def test_no_bootstrap_infos_is_noop(self):
+        r = _make_receiver(None)
+        r._send_abort_notification()  # must not raise
+        self.assertFalse(r.abort_notified)
+
+    def test_sends_to_every_peer(self):
+        socks = [_CapturingSock(), _CapturingSock()]
+        infos = [{"peer": 0}, {"peer": 1}]
+        r = _make_receiver(infos)
+        r._connect_to_bootstrap_server = lambda info: (socks[info["peer"]], _DummyLock())
+        r._send_abort_notification()
+        self.assertEqual(len(socks[0].sent), 1)
+        self.assertEqual(len(socks[1].sent), 1)
+
+
+class TestAbortWiring(unittest.TestCase):
+    def test_abort_notifies_and_marks_failed(self):
+        r = _make_receiver([{"peer": 1}])
+        called = []
+        r._send_abort_notification = lambda: called.append(True)
+        r.abort()
+        self.assertTrue(called)
+        self.assertEqual(r.conclude_state, KVPoll.Failed)
+        self.assertIn((777, KVPoll.Failed), r.kv_mgr.status_updates)
+
+
+class TestZmqWireFormat(unittest.TestCase):
+    def test_abort_and_ack_roundtrip(self):
+        import zmq
+
+        ctx = zmq.Context.instance()
+        pull = ctx.socket(zmq.PULL)
+        pull.bind("tcp://127.0.0.1:*")
+        ep = pull.getsockopt_string(zmq.LAST_ENDPOINT)
+        push = ctx.socket(zmq.PUSH)
+        push.connect(ep)
+        try:
+            push.send_multipart([b"ABORT", b"777", b"10.0.0.1", b"12345"])
+            parts = pull.recv_multipart()
+            self.assertEqual(parts[0], b"ABORT")
+            self.assertEqual(int(parts[1].decode()), 777)
+            self.assertEqual(parts[2].decode(), "10.0.0.1")
+            self.assertEqual(int(parts[3].decode()), 12345)
+
+            push.send_multipart([b"ABORT_ACK", b"777"])
+            ack = pull.recv_multipart()
+            self.assertEqual(ack[0], b"ABORT_ACK")
+            self.assertEqual(int(ack[1].decode()), 777)
+        finally:
+            push.close()
+            pull.close()
+
+
+if __name__ == "__main__":
+    unittest.main()