diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 3f6be17fb..8f36f6ee4 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/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( diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 3fbc37316..dba26671f 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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] = ( - decode_ip, - decode_port, + self.register_deferred_ack_target( + room_to_be_aborted, decode_ip, 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: - # Already completed/unknown: no in-flight write, ack now. + elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0: + # 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 ) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index b8cbe7925..ec642b84a 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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): - """Start a thread on the decode side to recv STAGING_REQ from prefill via ZMQ.""" + def _start_decode_listener_thread(self): + """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: - self._staging_outstanding.pop(room, None) - continue - - assert room in self.transfer_infos - - # Count each chunk once; the flag survives re-enqueue on defer. + # 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[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 - # release needs to wait for prefill-side NIXL transfer quiescence. + # Deferred KV release: register only after the status flip above (see + # 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 diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 9136395ed..0317832d7 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/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 diff --git a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py new file mode 100644 index 000000000..82d093bc0 --- /dev/null +++ b/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()