From a995e63118ff6dc5a9e0e66258f739110f980f74 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 17:54:44 +0800 Subject: [PATCH 01/10] [PD] Fix KV cache corruption on abort by notifying prefill and deferring slot release Signed-off-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 28 ++++++++ python/sglang/srt/disaggregation/decode.py | 72 ++++++++++++++++++- .../srt/disaggregation/mooncake/conn.py | 42 +++++++++++ python/sglang/srt/environ.py | 4 ++ 4 files changed, 143 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 94aa7d569ab8..e606ad3d96cf 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -877,6 +877,7 @@ def __init__( self.conclude_state: Optional[KVPoll] = None self.require_staging: bool = False self.init_time: Optional[float] = None + self.abort_initiated: 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) @@ -1073,12 +1074,39 @@ def clear(self) -> None: self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room, None) def abort(self): + self.abort_initiated = True self.kv_mgr.record_failure( self.bootstrap_room, "Aborted by AbortReq.", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.conclude_state = KVPoll.Failed + self._send_abort_notification() + + def _send_abort_notification(self): + if not hasattr(self, "bootstrap_infos") or self.bootstrap_infos is None: + return + 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/decode.py b/python/sglang/srt/disaggregation/decode.py index 022446bf4122..f1909a7dc458 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -244,6 +244,14 @@ def clear(self): self.mamba_pool.clear() +@dataclass +class DeferredRelease: + req: Req + bootstrap_room: int + metadata_buffer_index: int + enqueue_time: float + + @dataclass class DecodeRequest: req: Req @@ -1404,6 +1412,8 @@ def __init__( self.spec_algorithm = scheduler.spec_algorithm self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.staging_handler = None + self.deferred_releases: List[DeferredRelease] = [] + self.abort_grace_period = envs.SGLANG_DISAGGREGATION_ABORT_GRACE_PERIOD.get() def add(self, decode_req: DecodeRequest) -> None: self.queue.append(decode_req) @@ -1609,8 +1619,26 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) if self.scheduler.enable_hisparse: self.scheduler.hisparse_coordinator.request_finished(decode_req.req) - # release pre-allocated kv cache, but don't insert into the tree since it's failed - release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) + # Defer KV cache release for decode-initiated aborts to prevent + # RDMA corruption from in-flight prefill writes + is_decode_initiated_abort = getattr( + decode_req.kv_receiver, "decode_initiated_abort", False + ) + if is_decode_initiated_abort and self.abort_grace_period > 0: + self.deferred_releases.append( + DeferredRelease( + req=decode_req.req, + bootstrap_room=decode_req.req.bootstrap_room, + metadata_buffer_index=decode_req.metadata_buffer_index, + enqueue_time=time.monotonic(), + ) + ) + logger.debug( + f"Deferred KV release for room {decode_req.req.bootstrap_room}, " + f"waiting for ABORT_ACK or timeout" + ) + else: + release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) indices_to_remove.add(i) if self.scheduler.metrics_reporter.enable_metrics: self.scheduler.metrics_collector.increment_transfer_failed_reqs() @@ -1648,6 +1676,7 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req else: raise ValueError(f"Unexpected poll case: {poll}") + deferred_rooms = {d.bootstrap_room for d in self.deferred_releases} for i in indices_to_remove: if self.enable_staging and self.staging_handler.is_staging_room( self.queue[i].req.bootstrap_room @@ -1657,7 +1686,8 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) idx = self.queue[i].metadata_buffer_index assert idx != -1 - self.req_to_metadata_buffer_idx_allocator.free(idx) + if self.queue[i].req.bootstrap_room not in deferred_rooms: + self.req_to_metadata_buffer_idx_allocator.free(idx) self.queue = [ entry for i, entry in enumerate(self.queue) if i not in indices_to_remove @@ -1665,6 +1695,41 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req return transferred_reqs + def process_deferred_releases(self) -> None: + """Release deferred KV cache slots whose ABORT_ACK arrived or timeout expired.""" + if not self.deferred_releases: + return + + now = time.monotonic() + remaining = [] + prealloc_queue = self.scheduler.disagg_decode_prealloc_queue + kv_mgr = prealloc_queue.kv_manager if prealloc_queue else None + acked_rooms = getattr(kv_mgr, "_abort_acked_rooms", set()) if kv_mgr else set() + + for entry in self.deferred_releases: + acked = entry.bootstrap_room in acked_rooms + timed_out = (now - entry.enqueue_time) >= self.abort_grace_period + if acked or timed_out: + release_kv_cache(entry.req, self.tree_cache, is_insert=False) + self.req_to_metadata_buffer_idx_allocator.free( + entry.metadata_buffer_index + ) + if acked: + acked_rooms.discard(entry.bootstrap_room) + logger.debug( + f"Released deferred KV for room {entry.bootstrap_room} " + f"(ABORT_ACK received)" + ) + else: + logger.debug( + f"Released deferred KV for room {entry.bootstrap_room} " + f"(timeout after {self.abort_grace_period}s)" + ) + else: + remaining.append(entry) + + self.deferred_releases = remaining + class SchedulerDisaggregationDecodeMixin: @torch.no_grad() @@ -1878,6 +1943,7 @@ def process_decode_queue(self: Scheduler): transferred_reqs = ( self.disagg_decode_transfer_queue.pop_transferred() ) # the requests which kv has arrived + self.disagg_decode_transfer_queue.process_deferred_releases() if self.enable_hisparse: for req in transferred_reqs: # Direct-to-host: KV data already in host pool, skip staging diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 3d4406452e35..12505772a614 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -242,6 +242,7 @@ def __init__( self._init_staging_allocator() self._staging_handler = None self._chunk_writer_counts: dict = defaultdict(lambda: defaultdict(list)) + self._abort_acked_rooms: set = set() self.start_decode_thread() def init_engine(self): @@ -1386,6 +1387,40 @@ def bootstrap_thread(): 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")) + if room_to_be_aborted in self.request_status: + 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 unknown room {room_to_be_aborted}, " + f"ignoring" + ) + # 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 +1492,13 @@ def decode_thread(): self._handle_staging_req(msg) continue + # Prefill acknowledges abort notification + if msg[0] == b"ABORT_ACK": + ack_aborted_room = int(msg[1].decode("ascii")) + self._abort_acked_rooms.add(ack_aborted_room) + 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")) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8f4432c7d396..6c147554b8ed 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -303,6 +303,10 @@ class Envs: # max_num_reqs > 32). Increases pool capacity so more KV cache transfers # can overlap with decode execution without raising max_running_requests. SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS = EnvInt(0) + # Grace period (seconds) before releasing KV cache slots after a decode-initiated + # abort. The decode side holds resources until the prefill side ACKs or this timeout + # expires. Set to 0 to disable deferred release (immediate release on abort). + SGLANG_DISAGGREGATION_ABORT_GRACE_PERIOD = EnvFloat(5.0) # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. From a7986bbd86d01b97cb600aa1abde453f53770185 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 18:13:00 +0800 Subject: [PATCH 02/10] upd Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/decode.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f1909a7dc458..5140f39ad6d6 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1621,10 +1621,10 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req self.scheduler.hisparse_coordinator.request_finished(decode_req.req) # Defer KV cache release for decode-initiated aborts to prevent # RDMA corruption from in-flight prefill writes - is_decode_initiated_abort = getattr( - decode_req.kv_receiver, "decode_initiated_abort", False + is_abort_initiated = getattr( + decode_req.kv_receiver, "abort_initiated", False ) - if is_decode_initiated_abort and self.abort_grace_period > 0: + if is_abort_initiated and self.abort_grace_period > 0: self.deferred_releases.append( DeferredRelease( req=decode_req.req, From 3465c2b3425844575d6ecbf0cad19ef91f1c48b5 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 19:05:34 +0800 Subject: [PATCH 03/10] upd Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/decode.py | 4 ++++ python/sglang/srt/disaggregation/mooncake/conn.py | 13 ++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5140f39ad6d6..790a5f80b1ba 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1729,6 +1729,10 @@ def process_deferred_releases(self) -> None: remaining.append(entry) self.deferred_releases = remaining + # Prune late ACKs" that arrived after the request was already aborted + if acked_rooms: + active_rooms = {d.bootstrap_room for d in self.deferred_releases} + acked_rooms -= acked_rooms - active_rooms class SchedulerDisaggregationDecodeMixin: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 12505772a614..7a6130fe1dd6 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1162,6 +1162,13 @@ def transfer_worker( while True: try: kv_chunk: TransferKVChunk = queue.get() + + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Failed + ): + continue + if self.enable_trace: kv_chunk.trace_ctx.rebuild_thread_context() kv_chunk.trace_ctx.trace_slice_start( @@ -1392,7 +1399,11 @@ def bootstrap_thread(): 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: + # 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}, " From 08968ae5b463e783fb1941d4ea51472397fe140c Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 19:29:41 +0800 Subject: [PATCH 04/10] upd Signed-off-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 10 +++++++--- python/sglang/srt/disaggregation/decode.py | 6 +++--- .../srt/disaggregation/mooncake/conn.py | 19 ++++++++++++------- 3 files changed, 22 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index e606ad3d96cf..ea616e31abc2 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -877,7 +877,7 @@ def __init__( self.conclude_state: Optional[KVPoll] = None self.require_staging: bool = False self.init_time: Optional[float] = None - self.abort_initiated: bool = False + 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) @@ -1063,6 +1063,9 @@ def _check_waiting_timeout(self) -> Optional[KVPoll]: f"in KVPoll.WaitingForInput", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + if not self.abort_notified: + self._send_abort_notification() + self.abort_notified = True return KVPoll.Failed def failure_exception(self): @@ -1074,14 +1077,15 @@ def clear(self) -> None: self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room, None) def abort(self): - self.abort_initiated = True self.kv_mgr.record_failure( self.bootstrap_room, "Aborted by AbortReq.", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.conclude_state = KVPoll.Failed - self._send_abort_notification() + if not self.abort_notified: + self._send_abort_notification() + self.abort_notified = True def _send_abort_notification(self): if not hasattr(self, "bootstrap_infos") or self.bootstrap_infos is None: diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 790a5f80b1ba..fa809d26d8b6 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1621,10 +1621,10 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req self.scheduler.hisparse_coordinator.request_finished(decode_req.req) # Defer KV cache release for decode-initiated aborts to prevent # RDMA corruption from in-flight prefill writes - is_abort_initiated = getattr( - decode_req.kv_receiver, "abort_initiated", False + is_abort_notified = getattr( + decode_req.kv_receiver, "abort_notified", False ) - if is_abort_initiated and self.abort_grace_period > 0: + if is_abort_notified and self.abort_grace_period > 0: self.deferred_releases.append( DeferredRelease( req=decode_req.req, diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 7a6130fe1dd6..45376cc1f8e1 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1162,13 +1162,6 @@ def transfer_worker( while True: try: kv_chunk: TransferKVChunk = queue.get() - - if ( - kv_chunk.room not in self.request_status - or self.check_status(kv_chunk.room) == KVPoll.Failed - ): - continue - if self.enable_trace: kv_chunk.trace_ctx.rebuild_thread_context() kv_chunk.trace_ctx.trace_slice_start( @@ -1176,6 +1169,18 @@ def transfer_worker( MooncakeRequestStage.MOONCAKE_WORKER_SEND.level, ) + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Failed + ): + 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 From 87166d61b621d6569c812228b6de5f01c37a53ec Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 19:52:57 +0800 Subject: [PATCH 05/10] upd Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/common/conn.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index ea616e31abc2..b4ed4f0dbc30 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -1063,7 +1063,11 @@ def _check_waiting_timeout(self) -> Optional[KVPoll]: f"in KVPoll.WaitingForInput", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) - if not self.abort_notified: + 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 @@ -1083,13 +1087,15 @@ def abort(self): ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.conclude_state = KVPoll.Failed - if not self.abort_notified: + 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): - if not hasattr(self, "bootstrap_infos") or self.bootstrap_infos is None: - return for bootstrap_info in self.bootstrap_infos: # Best-effort notification to prefill side that this request was aborted. try: From 03f7baf369ec0ab89fb69b2a8572397c02429c3d Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 20:08:23 +0800 Subject: [PATCH 06/10] revert deffer Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/decode.py | 75 +--------------------- python/sglang/srt/environ.py | 4 -- 2 files changed, 2 insertions(+), 77 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index fa809d26d8b6..68400d145820 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -244,14 +244,6 @@ def clear(self): self.mamba_pool.clear() -@dataclass -class DeferredRelease: - req: Req - bootstrap_room: int - metadata_buffer_index: int - enqueue_time: float - - @dataclass class DecodeRequest: req: Req @@ -1412,8 +1404,6 @@ def __init__( self.spec_algorithm = scheduler.spec_algorithm self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.staging_handler = None - self.deferred_releases: List[DeferredRelease] = [] - self.abort_grace_period = envs.SGLANG_DISAGGREGATION_ABORT_GRACE_PERIOD.get() def add(self, decode_req: DecodeRequest) -> None: self.queue.append(decode_req) @@ -1619,26 +1609,8 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) if self.scheduler.enable_hisparse: self.scheduler.hisparse_coordinator.request_finished(decode_req.req) - # Defer KV cache release for decode-initiated aborts to prevent - # RDMA corruption from in-flight prefill writes - is_abort_notified = getattr( - decode_req.kv_receiver, "abort_notified", False - ) - if is_abort_notified and self.abort_grace_period > 0: - self.deferred_releases.append( - DeferredRelease( - req=decode_req.req, - bootstrap_room=decode_req.req.bootstrap_room, - metadata_buffer_index=decode_req.metadata_buffer_index, - enqueue_time=time.monotonic(), - ) - ) - logger.debug( - f"Deferred KV release for room {decode_req.req.bootstrap_room}, " - f"waiting for ABORT_ACK or timeout" - ) - else: - release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) + # release pre-allocated kv cache, but don't insert into the tree since it's failed + release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) indices_to_remove.add(i) if self.scheduler.metrics_reporter.enable_metrics: self.scheduler.metrics_collector.increment_transfer_failed_reqs() @@ -1676,7 +1648,6 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req else: raise ValueError(f"Unexpected poll case: {poll}") - deferred_rooms = {d.bootstrap_room for d in self.deferred_releases} for i in indices_to_remove: if self.enable_staging and self.staging_handler.is_staging_room( self.queue[i].req.bootstrap_room @@ -1686,8 +1657,6 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) idx = self.queue[i].metadata_buffer_index assert idx != -1 - if self.queue[i].req.bootstrap_room not in deferred_rooms: - self.req_to_metadata_buffer_idx_allocator.free(idx) self.queue = [ entry for i, entry in enumerate(self.queue) if i not in indices_to_remove @@ -1695,45 +1664,6 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req return transferred_reqs - def process_deferred_releases(self) -> None: - """Release deferred KV cache slots whose ABORT_ACK arrived or timeout expired.""" - if not self.deferred_releases: - return - - now = time.monotonic() - remaining = [] - prealloc_queue = self.scheduler.disagg_decode_prealloc_queue - kv_mgr = prealloc_queue.kv_manager if prealloc_queue else None - acked_rooms = getattr(kv_mgr, "_abort_acked_rooms", set()) if kv_mgr else set() - - for entry in self.deferred_releases: - acked = entry.bootstrap_room in acked_rooms - timed_out = (now - entry.enqueue_time) >= self.abort_grace_period - if acked or timed_out: - release_kv_cache(entry.req, self.tree_cache, is_insert=False) - self.req_to_metadata_buffer_idx_allocator.free( - entry.metadata_buffer_index - ) - if acked: - acked_rooms.discard(entry.bootstrap_room) - logger.debug( - f"Released deferred KV for room {entry.bootstrap_room} " - f"(ABORT_ACK received)" - ) - else: - logger.debug( - f"Released deferred KV for room {entry.bootstrap_room} " - f"(timeout after {self.abort_grace_period}s)" - ) - else: - remaining.append(entry) - - self.deferred_releases = remaining - # Prune late ACKs" that arrived after the request was already aborted - if acked_rooms: - active_rooms = {d.bootstrap_room for d in self.deferred_releases} - acked_rooms -= acked_rooms - active_rooms - class SchedulerDisaggregationDecodeMixin: @torch.no_grad() @@ -1947,7 +1877,6 @@ def process_decode_queue(self: Scheduler): transferred_reqs = ( self.disagg_decode_transfer_queue.pop_transferred() ) # the requests which kv has arrived - self.disagg_decode_transfer_queue.process_deferred_releases() if self.enable_hisparse: for req in transferred_reqs: # Direct-to-host: KV data already in host pool, skip staging diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6c147554b8ed..8f4432c7d396 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -303,10 +303,6 @@ class Envs: # max_num_reqs > 32). Increases pool capacity so more KV cache transfers # can overlap with decode execution without raising max_running_requests. SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS = EnvInt(0) - # Grace period (seconds) before releasing KV cache slots after a decode-initiated - # abort. The decode side holds resources until the prefill side ACKs or this timeout - # expires. Set to 0 to disable deferred release (immediate release on abort). - SGLANG_DISAGGREGATION_ABORT_GRACE_PERIOD = EnvFloat(5.0) # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. From cd39d94eb99eb240798485b94704b845d33ca5b7 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 20:09:13 +0800 Subject: [PATCH 07/10] fix Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/decode.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 68400d145820..022446bf4122 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1657,6 +1657,7 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) idx = self.queue[i].metadata_buffer_index assert idx != -1 + self.req_to_metadata_buffer_idx_allocator.free(idx) self.queue = [ entry for i, entry in enumerate(self.queue) if i not in indices_to_remove From dd814f505c1608d2928afd5a5ca78a4a1009f2c5 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 20:12:16 +0800 Subject: [PATCH 08/10] clean Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/mooncake/conn.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 45376cc1f8e1..296ca71ebf01 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -242,7 +242,6 @@ def __init__( self._init_staging_allocator() self._staging_handler = None self._chunk_writer_counts: dict = defaultdict(lambda: defaultdict(list)) - self._abort_acked_rooms: set = set() self.start_decode_thread() def init_engine(self): @@ -1511,7 +1510,6 @@ def decode_thread(): # Prefill acknowledges abort notification if msg[0] == b"ABORT_ACK": ack_aborted_room = int(msg[1].decode("ascii")) - self._abort_acked_rooms.add(ack_aborted_room) logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}") continue From a3162a4238ca0305fe14c562e462136f44d215d6 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 20:20:03 +0800 Subject: [PATCH 09/10] add log Signed-off-by: Shangming Cai --- python/sglang/srt/disaggregation/mooncake/conn.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 296ca71ebf01..6fb5ed1af6a6 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1172,6 +1172,9 @@ def transfer_worker( 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, From e3e8c3c3372f9893ccef6bda7575bb6054f5d1e4 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Fri, 5 Jun 2026 23:26:24 +0800 Subject: [PATCH 10/10] update debug log --- python/sglang/srt/disaggregation/mooncake/conn.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 6fb5ed1af6a6..ccae5cc415c4 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1418,8 +1418,8 @@ def bootstrap_thread(): ) else: logger.debug( - f"Received abort notification for unknown room {room_to_be_aborted}, " - f"ignoring" + f"Received abort notification for room {room_to_be_aborted}, " + f"ignoring (already completed or unknown)" ) # Send ACK back to decode endpoint try: @@ -1512,6 +1512,7 @@ def decode_thread(): # 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