[PD] Deferred decode-side KV release for the NIXL backend (#35360)

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

@@ -236,6 +236,9 @@ class CommonKVManager(BaseKVManager):
             )
             self.register_to_bootstrap()
             self.transfer_infos = {}
+            # Deferred KV release: aborted room -> (decode_ip, decode_port);
+            # ack held until the transfer drains.
+            self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {}
             self.req_to_decode_prefix_len: Dict[int, int] = {}
             self.decode_kv_args_table = {}
             self.pp_group = get_pp_group()
@@ -364,6 +367,47 @@ class CommonKVManager(BaseKVManager):
     def clear_deferred_abort_state(self, bootstrap_room: int) -> None:
         self._deferred_abort_ack_tracker.pop(bootstrap_room, None)
 
+    def _prefill_unique_rank(self) -> int:
+        """Stable per-sender id, matching what the transfer worker syncs on Success."""
+        return (
+            self.attn_tp_rank * (self.pp_size * self.attn_cp_size)
+            + self.pp_rank * self.attn_cp_size
+            + self.attn_cp_rank
+        )
+
+    def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None:
+        """Best-effort ack that this rank's transfer for an aborted room drained."""
+        try:
+            na = NetworkAddress(decode_ip, decode_port)
+            self._send_multipart_locked(
+                na.to_tcp(),
+                [
+                    b"ABORT_ACK",
+                    str(room).encode("ascii"),
+                    str(self._prefill_unique_rank()).encode("ascii"),
+                ],
+                is_ipv6=na.is_ipv6,
+            )
+        except Exception as e:
+            logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}")
+
+    def _maybe_ack_drained_abort(self, room: int) -> None:
+        """Send the deferred ack once an aborted room's chunks have drained
+        (outstanding == 0). pop() makes it fire at most once."""
+        if self._staging_outstanding.get(room, 0) > 0:
+            return
+        target = self._deferred_ack_targets.pop(room, None)
+        if target is not None:
+            self._send_abort_ack(target[0], target[1], room)
+
+    def register_deferred_ack_target(
+        self, room: int, decode_ip: str, decode_port: int
+    ) -> None:
+        """Hold this room's ack until its transfer drains. Callers must mark the
+        room Failed FIRST -- registering while it still accepts chunks lets the
+        worker ack, then a new chunk writes pages the decode already released."""
+        self._deferred_ack_targets[room] = (decode_ip, decode_port)
+
     def get_kv_replica_factor(self) -> int:
         if self._kv_replica_factor is None:
             logger.warning_once(

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

@@ -8,7 +8,7 @@ import struct
 import threading
 import time
 from collections import defaultdict
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import List, Optional, Set, Tuple, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -214,10 +214,6 @@ class MooncakeKVManager(CommonKVManager):
             # Per-room count of chunks not yet transferred; teardown waits for
             # zero so a deferred chunk is not dropped by an early conclude.
             self._staging_outstanding = defaultdict(int)
-            # Deferred KV release: aborted room -> (decode_ip, decode_port), ack
-            # held until the transfer drains. Written by the bootstrap thread,
-            # popped by the single transfer worker that owns the room.
-            self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {}
             self.session_lock = threading.Lock()
             # Determine the number of threads to use for kv sender
             cpu_count = os.cpu_count()
@@ -1616,39 +1612,6 @@ class MooncakeKVManager(CommonKVManager):
             is_ipv6=na.is_ipv6,
         )
 
-    def _prefill_unique_rank(self) -> int:
-        """Stable per-sender id, matching what the transfer worker syncs on Success."""
-        return (
-            self.attn_tp_rank * (self.pp_size * self.attn_cp_size)
-            + self.pp_rank * self.attn_cp_size
-            + self.attn_cp_rank
-        )
-
-    def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None:
-        """Best-effort ack that this rank's transfer for an aborted room drained."""
-        try:
-            na = NetworkAddress(decode_ip, decode_port)
-            self._send_multipart_locked(
-                na.to_tcp(),
-                [
-                    b"ABORT_ACK",
-                    str(room).encode("ascii"),
-                    str(self._prefill_unique_rank()).encode("ascii"),
-                ],
-                is_ipv6=na.is_ipv6,
-            )
-        except Exception as e:
-            logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}")
-
-    def _maybe_ack_drained_abort(self, room: int) -> None:
-        """Send the deferred ack once an aborted room's chunks have drained
-        (outstanding == 0). pop() makes it fire at most once."""
-        if self._staging_outstanding.get(room, 0) > 0:
-            return
-        target = self._deferred_ack_targets.pop(room, None)
-        if target is not None:
-            self._send_abort_ack(target[0], target[1], room)
-
     def transfer_worker(
         self,
         queue: FastQueue,
@@ -1674,6 +1637,14 @@ class MooncakeKVManager(CommonKVManager):
                         MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
                     )
 
+                # Counted at dequeue, before the status check, so
+                # `outstanding == 0` means nothing is dequeued or in flight --
+                # the predicate the abort ack relies on. The flag survives
+                # re-enqueue on defer.
+                if not kv_chunk.staging_counted:
+                    self._staging_outstanding[kv_chunk.room] += 1
+                    kv_chunk.staging_counted = True
+
                 if (
                     kv_chunk.room not in self.request_status
                     or self.check_status(kv_chunk.room) == KVPoll.Failed
@@ -1693,11 +1664,6 @@ class MooncakeKVManager(CommonKVManager):
                         self._maybe_ack_drained_abort(kv_chunk.room)
                     continue
 
-                # Count each chunk once; the flag survives re-enqueue on defer.
-                if not kv_chunk.staging_counted:
-                    self._staging_outstanding[kv_chunk.room] += 1
-                    kv_chunk.staging_counted = True
-
                 if (
                     self.enable_staging
                     and staging_strategy is None
@@ -2042,16 +2008,20 @@ class MooncakeKVManager(CommonKVManager):
                         # flight, decode falls back to the release timeout.
                         if room_active:
                             self.update_status(room_to_be_aborted, KVPoll.Failed)
-                            self._deferred_ack_targets[room_to_be_aborted] = (
+                            self.register_deferred_ack_target(
-                                decode_ip,
+                                room_to_be_aborted, decode_ip, decode_port
-                                decode_port,
                             )
+                            # Try once: the room may already be quiescent and
+                            # never revisited by the worker.
+                            self._maybe_ack_drained_abort(room_to_be_aborted)
                             logger.debug(
                                 f"Received abort notification for room {room_to_be_aborted}, "
                                 f"marked as Failed; ACK deferred until transfer drains"
                             )
-                        else:
+                        elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0:
-                            # Already completed/unknown: no in-flight write, ack now.
+                            # Concluded/unknown AND quiescent: ack now. A cleared
+                            # room is not automatically quiescent -- clear() can
+                            # drop a room whose chunk is still transferring.
                             self._send_abort_ack(
                                 decode_ip, decode_port, room_to_be_aborted
                             )

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

@@ -512,7 +512,8 @@ class NixlKVManager(CommonKVManager):
             if self.enable_staging:
                 self._init_staging_decode_ctx()
                 self._staging_handler = None
-                self._start_decode_staging_thread()
+            if self.enable_staging or self.enable_deferred_decode_kv_release:
+                self._start_decode_listener_thread()
             self._start_heartbeat_checker_thread()
         else:
             raise ValueError(
@@ -592,21 +593,30 @@ class NixlKVManager(CommonKVManager):
 
         return is_watermark_ready(self._staging_ctx, agent_name, alloc_round, alloc_end)
 
-    def _start_decode_staging_thread(self):
+    def _start_decode_listener_thread(self):
-        """Start a thread on the decode side to recv STAGING_REQ from prefill via ZMQ."""
+        """Decode-side ZMQ listener for STAGING_REQ and ABORT_ACK. A thread, not
+        NIXL notifs: the decode agent has no progress thread, so notifs only drain
+        inside a live receiver's poll() and would be missed while idle."""
 
-        def decode_staging_thread():
+        def decode_listener_thread():
             while True:
                 msg = self.server_socket.recv_multipart()
                 if msg[0] == b"STAGING_REQ":
                     self._handle_staging_req(msg)
                     continue
+                if msg[0] == b"ABORT_ACK":
+                    # Drain ack for an aborted room; aggregate per prefill rank.
+                    if len(msg) >= 3:
+                        self.note_abort_ack(
+                            int(msg[1].decode("ascii")), int(msg[2].decode("ascii"))
+                        )
+                    continue
                 logger.warning(
-                    "decode_staging_thread: unexpected message tag %s",
+                    "decode_listener_thread: unexpected message tag %s",
                     msg[0][:20],
                 )
 
-        threading.Thread(target=decode_staging_thread, daemon=True).start()
+        threading.Thread(target=decode_listener_thread, daemon=True).start()
 
     def _handle_staging_req(self, msg):
         from sglang.srt.disaggregation.common.staging_handler import (
@@ -1119,17 +1129,23 @@ class NixlKVManager(CommonKVManager):
             room = kv_chunk.room
             handles: List[Any] = []
             try:
-                if self.check_status(room) == KVPoll.Failed:
+                # Counted at dequeue, before the status check, so
-                    self._staging_outstanding.pop(room, None)
+                # `outstanding == 0` means nothing is dequeued or in flight --
-                    continue
+                # the predicate the abort ack relies on. The flag survives
-
+                # re-enqueue on defer.
-                assert room in self.transfer_infos
-
-                # Count each chunk once; the flag survives re-enqueue on defer.
                 if not kv_chunk.staging_counted:
                     self._staging_outstanding[room] += 1
                     kv_chunk.staging_counted = True
 
+                if self.check_status(room) == KVPoll.Failed:
+                    self._staging_outstanding.pop(room, None)
+                    if self.enable_deferred_decode_kv_release:
+                        # Skipped => nothing written for this aborted room; ack.
+                        self._maybe_ack_drained_abort(room)
+                    continue
+
+                assert room in self.transfer_infos
+
                 # Lazily build a per-worker staging strategy bound to this
                 # worker's private staging buffer (matches mooncake).
                 if (
@@ -1344,6 +1360,10 @@ class NixlKVManager(CommonKVManager):
                     time.sleep(0)
 
                 self._staging_outstanding[room] -= 1
+                if self.enable_deferred_decode_kv_release:
+                    # Handles all DONE => this room's writes landed; ack if it
+                    # was aborted and nothing else is outstanding.
+                    self._maybe_ack_drained_abort(room)
                 if kv_chunk.is_last_chunk:
                     self.update_status(room, KVPoll.Success)
                 elif self.check_status(room) != KVPoll.Success:
@@ -1383,6 +1403,8 @@ class NixlKVManager(CommonKVManager):
                 self.exceptions[room] = e
                 self.record_failure(room, str(e))
                 self.update_status(room, KVPoll.Failed)
+                # No ack here on purpose: the DONE barrier bails on the first
+                # ERR, so siblings may still be writing; fall back to the timeout.
 
     def register_buffer_to_engine(self):
         self.kv_descs = []
@@ -2628,14 +2650,17 @@ class NixlKVManager(CommonKVManager):
 
         try:
             room_to_be_aborted = int(msg[1].decode("ascii"))
+            decode_ip = msg[2].decode("ascii") if len(msg) > 2 else None
+            decode_port = int(msg[3].decode("ascii")) if len(msg) > 3 else None
         except Exception as e:
             logger.debug(f"Ignoring malformed abort notification: {e}")
             return True
 
-        if (
+        room_active = (
             room_to_be_aborted in self.request_status
             and self.check_status(room_to_be_aborted) != KVPoll.Success
-        ):
+        )
+        if room_active:
             self.record_failure(
                 room_to_be_aborted,
                 "Aborted by decode-side abort notification.",
@@ -2651,8 +2676,20 @@ class NixlKVManager(CommonKVManager):
                 f"ignoring (already completed or unknown)"
             )
 
-        # TODO: Define real ACK/deferred-release semantics if decode-side buffer
+        # Deferred KV release: register only after the status flip above (see
-        # release needs to wait for prefill-side NIXL transfer quiescence.
+        # register_deferred_ack_target), then try once -- the room may already be
+        # quiescent and never revisited by the worker. A concluded/unknown room is
+        # acked only when nothing is still counted for it: the ERR path abandons
+        # sibling handles that may still be writing and clear() then drops the
+        # room, so "unknown" alone does not imply quiescent.
+        if self.enable_deferred_decode_kv_release and decode_port is not None:
+            if room_active:
+                self.register_deferred_ack_target(
+                    room_to_be_aborted, decode_ip, decode_port
+                )
+                self._maybe_ack_drained_abort(room_to_be_aborted)
+            elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0:
+                self._send_abort_ack(decode_ip, decode_port, room_to_be_aborted)
 
         return True
 

test/registered/unit/disaggregation/test_nixl_backend_basic.py

@@ -347,6 +347,9 @@ class TestNixlAbortHandling(CustomTestCase):
         mgr._connect = MagicMock()
         mgr.failure_lock = threading.Lock()
         mgr.failure_records = {}
+        # These cases cover the legacy no-ack behavior; the deferred-release ack
+        # path is exercised in test_nixl_deferred_kv_release.py.
+        mgr.enable_deferred_decode_kv_release = False
         return mgr
 
     def test_given_known_incomplete_room_when_abort_arrives_then_room_fails_without_ack(
@@ -466,6 +469,7 @@ class TestNixlTransferWorker(CustomTestCase):
         }
         mgr.req_to_decode_prefix_len = {room: 4}
         mgr.enable_staging = False
+        mgr.enable_deferred_decode_kv_release = False
         mgr._staging_ctx = None
         mgr._staging_outstanding = defaultdict(int)
         mgr.is_mla_backend = False

test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py

@@ -0,0 +1,188 @@
+"""Deferred decode-side KV release on the NIXL backend.
+
+When a decode request is aborted while its prefill->decode transfer may still be
+in flight, the decode holds its KV pages until every prefill rank acks that its
+transfer drained. NIXL transfers are asynchronous (agent.transfer() posts, the
+worker polls check_xfer_state), so the ack must come from the transfer worker
+after its DONE barrier -- never from the bootstrap thread for an active room.
+"""
+
+import unittest
+from unittest.mock import MagicMock
+
+from sglang.srt.disaggregation.base.conn import KVPoll
+from sglang.srt.disaggregation.common.conn import CommonKVManager
+from sglang.srt.disaggregation.nixl.conn import NixlKVManager
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _prefill_mgr(cls=CommonKVManager, enabled=True):
+    """Bare manager carrying only the prefill-side deferred-ack state."""
+    mgr = cls.__new__(cls)
+    mgr.enable_deferred_decode_kv_release = enabled
+    mgr._deferred_ack_targets = {}
+    mgr._staging_outstanding = {}
+    mgr.request_status = {}
+    mgr._sent = []
+    # Capture acks instead of opening a socket.
+    mgr._send_abort_ack = lambda ip, port, room: mgr._sent.append((ip, port, room))
+    return mgr
+
+
+class TestDeferredAckTargets(CustomTestCase):
+    def test_ack_held_until_outstanding_drains(self):
+        mgr = _prefill_mgr()
+        mgr.register_deferred_ack_target(7, "10.0.0.1", 5000)
+
+        mgr._staging_outstanding[7] = 1
+        mgr._maybe_ack_drained_abort(7)
+        self.assertEqual(mgr._sent, [])  # still writing -> no ack
+
+        mgr._staging_outstanding[7] = 0
+        mgr._maybe_ack_drained_abort(7)
+        self.assertEqual(mgr._sent, [("10.0.0.1", 5000, 7)])
+
+    def test_ack_fires_at_most_once(self):
+        mgr = _prefill_mgr()
+        mgr.register_deferred_ack_target(8, "10.0.0.2", 5001)
+        mgr._maybe_ack_drained_abort(8)
+        mgr._maybe_ack_drained_abort(8)
+        self.assertEqual(len(mgr._sent), 1)
+        self.assertNotIn(8, mgr._deferred_ack_targets)
+
+    def test_unregistered_room_is_noop(self):
+        mgr = _prefill_mgr()
+        mgr._maybe_ack_drained_abort(999)
+        self.assertEqual(mgr._sent, [])
+
+    def test_prefill_unique_rank_matches_success_sync_formula(self):
+        mgr = CommonKVManager.__new__(CommonKVManager)
+        mgr.attn_tp_rank, mgr.pp_size, mgr.attn_cp_size = 2, 3, 4
+        mgr.pp_rank, mgr.attn_cp_rank = 1, 3
+        self.assertEqual(mgr._prefill_unique_rank(), 2 * (3 * 4) + 1 * 4 + 3)
+
+
+class TestNixlAbortNotification(CustomTestCase):
+    """_handle_abort_notification is the prefill bootstrap-thread entry point."""
+
+    @staticmethod
+    def _abort_msg(room=11, ip="10.0.0.3", port=6000):
+        return [
+            b"ABORT",
+            str(room).encode("ascii"),
+            ip.encode("ascii"),
+            str(port).encode("ascii"),
+        ]
+
+    def _mgr(self, enabled=True, room=11, status=KVPoll.WaitingForInput):
+        mgr = _prefill_mgr(NixlKVManager, enabled=enabled)
+        if status is not None:
+            mgr.request_status[room] = status
+        mgr.record_failure = MagicMock()
+        mgr.update_status = MagicMock(
+            side_effect=lambda r, s: mgr.request_status.__setitem__(r, s)
+        )
+        mgr.check_status = lambda r: mgr.request_status[r]
+        return mgr
+
+    def test_in_flight_room_registers_target_and_does_not_ack_yet(self):
+        # A counted chunk holds the ack: only the worker knows when it landed.
+        mgr = self._mgr()
+        mgr._staging_outstanding[11] = 1
+        self.assertTrue(mgr._handle_abort_notification(self._abort_msg()))
+
+        self.assertEqual(mgr._deferred_ack_targets[11], ("10.0.0.3", 6000))
+        self.assertEqual(mgr._sent, [])
+        # Marked Failed first, so no new chunk can be enqueued for the room.
+        self.assertEqual(mgr.request_status[11], KVPoll.Failed)
+
+    def test_quiescent_active_room_acks_without_waiting_for_a_worker_visit(self):
+        # Window 2: chunks already drained with none left to come, so the worker
+        # never revisits the room -- acking here keeps it off the timeout path.
+        mgr = self._mgr()
+        self.assertTrue(mgr._handle_abort_notification(self._abort_msg()))
+
+        self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)])
+        self.assertEqual(mgr._deferred_ack_targets, {})
+
+    def test_worker_skip_before_registration_still_acks(self):
+        # Window 1: the worker can pass its skip point between the Failed flip
+        # and registration; the ack attempt at registration covers that.
+        mgr = self._mgr()
+        mgr._staging_outstanding[11] = 1
+
+        real_update = mgr.update_status.side_effect
+
+        def failed_then_worker_skips(room, status):
+            real_update(room, status)
+            # Worker dequeues, sees Failed, uncounts, and finds no target yet.
+            mgr._staging_outstanding.pop(room, None)
+            mgr._maybe_ack_drained_abort(room)
+
+        mgr.update_status = MagicMock(side_effect=failed_then_worker_skips)
+        self.assertTrue(mgr._handle_abort_notification(self._abort_msg()))
+
+        self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)])
+        self.assertEqual(mgr._deferred_ack_targets, {})
+
+    def test_concluded_room_acks_immediately(self):
+        # Concluded and quiescent: ack straight away.
+        mgr = self._mgr(status=None)
+        mgr.check_status = lambda r: KVPoll.Success
+        self.assertTrue(mgr._handle_abort_notification(self._abort_msg()))
+
+        self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)])
+        self.assertEqual(mgr._deferred_ack_targets, {})
+
+    def test_cleared_room_with_outstanding_chunk_does_not_ack(self):
+        # The ERR path abandons sibling handles that may still be writing and
+        # leaves the chunk counted; clear() then drops the room. Acking on
+        # "unknown room" alone would release decode pages under those writes.
+        mgr = self._mgr(status=None)  # room absent == cleared/unknown
+        mgr._staging_outstanding[11] = 1
+        self.assertTrue(mgr._handle_abort_notification(self._abort_msg()))
+
+        self.assertEqual(mgr._sent, [])
+        self.assertEqual(mgr._deferred_ack_targets, {})
+
+    def test_feature_off_registers_nothing_and_acks_nothing(self):
+        mgr = self._mgr(enabled=False)
+        self.assertTrue(mgr._handle_abort_notification(self._abort_msg()))
+
+        self.assertEqual(mgr._deferred_ack_targets, {})
+        self.assertEqual(mgr._sent, [])
+        # Legacy behavior preserved: the room is still failed.
+        self.assertEqual(mgr.request_status[11], KVPoll.Failed)
+
+    def test_legacy_two_frame_abort_is_tolerated(self):
+        # Older peers send [ABORT, room] with no return address.
+        mgr = self._mgr()
+        self.assertTrue(mgr._handle_abort_notification([b"ABORT", b"11"]))
+        self.assertEqual(mgr._deferred_ack_targets, {})
+        self.assertEqual(mgr._sent, [])
+
+    def test_non_abort_message_is_not_claimed(self):
+        mgr = self._mgr()
+        self.assertFalse(mgr._handle_abort_notification([b"STAGING_REQ", b"11"]))
+
+
+class TestNixlDecodeAckIngest(CustomTestCase):
+    def test_abort_ack_is_aggregated_per_rank(self):
+        # Mirrors the decode listener thread's ABORT_ACK branch.
+        mgr = CommonKVManager.__new__(CommonKVManager)
+        mgr._deferred_abort_ack_tracker = {}
+        mgr.register_deferred_abort_room(21)
+
+        for rank in (b"0", b"1", b"1"):
+            msg = [b"ABORT_ACK", b"21", rank]
+            mgr.note_abort_ack(int(msg[1].decode()), int(msg[2].decode()))
+
+        self.assertFalse(mgr.is_abort_release_safe(21, required_acks=3))
+        self.assertTrue(mgr.is_abort_release_safe(21, required_acks=2))
+
+
+if __name__ == "__main__":
+    unittest.main()