diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 3c916229aebc..da80d6d2c1a9 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -159,6 +159,7 @@ def __init__( self.session_pool_lock = threading.Lock() self.addr_to_rooms_tracker: Dict[str, Set[int]] = defaultdict(set) self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set) + self.prefill_abort_tracker: Dict[int, Set[int]] = defaultdict(set) # Heartbeat interval should be at least 2 seconds self.heartbeat_interval = max( envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0 @@ -768,6 +769,7 @@ def _setup_routes(self): self.app.router.add_route("*", "/route", self._handle_route) self.app.router.add_post("/register_dp_rank", self._handle_register_dp_rank) self.app.router.add_post("/query_dp_ranks", self._handle_query_dp_ranks) + self.app.router.add_post("/abort_room", self._handle_abort_room) self.app.router.add_get("/health", self._handle_health_check) async def _handle_health_check(self, request): @@ -936,6 +938,50 @@ async def _handle_query_dp_ranks(self, request: web.Request): result[str(room_int)] = self.room_to_dp_rank[room_int]["dp_rank"] return web.json_response(result, status=200) + async def _handle_abort_room(self, request: web.Request): + data = await request.json() + bootstrap_room = int(data["bootstrap_room"]) + endpoint = str(data["endpoint"]) + dst_port = int(data["dst_port"]) + prefill_dp_rank = int(data["prefill_dp_rank"]) + target_tp_ranks = [int(x) for x in data["target_tp_ranks"]] + target_cp_ranks = [int(x) for x in data["target_cp_ranks"]] + target_pp_ranks = [int(x) for x in data["target_pp_ranks"]] + + payload = [ + b"ABORT_REQ", + str(bootstrap_room).encode("ascii"), + endpoint.encode("ascii"), + str(dst_port).encode("ascii"), + ] + fanout_targets = [] + async with self.lock: + dp_group = self.prefill_port_table.get(prefill_dp_rank, {}) + for cp_rank in target_cp_ranks: + cp_group = dp_group.get(cp_rank, {}) + for tp_rank in target_tp_ranks: + tp_group = cp_group.get(tp_rank, {}) + for pp_rank in target_pp_ranks: + rank_info = tp_group.get(pp_rank) + if rank_info is not None: + fanout_targets.append( + (rank_info.rank_ip, rank_info.rank_port) + ) + + for rank_ip, rank_port in fanout_targets: + na = NetworkAddress(rank_ip, rank_port) + sock, lock = CommonKVReceiver._connect(na.to_tcp(), is_ipv6=na.is_ipv6) + with lock: + sock.send_multipart(payload) + + logger.info( + "[PD_ABORT] bootstrap fanout room=%s prefill_dp_rank=%s targets=%s", + bootstrap_room, + prefill_dp_rank, + fanout_targets, + ) + return web.Response(text="OK", status=200) + async def _cleanup_expired_entries(self): """Remove entries older than cleanup interval from room_to_dp_rank.""" while True: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index fe0d651038f3..f665e0ec42ae 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -13,6 +13,7 @@ import numpy as np import numpy.typing as npt +import requests from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll from sglang.srt.disaggregation.common.conn import ( @@ -190,6 +191,8 @@ def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data): class MooncakeKVManager(CommonKVManager): AUX_DATA_HEADER = b"AUX_DATA" + ABORT_REQ_HEADER = b"ABORT_REQ" + ABORT_ACK_HEADER = b"ABORT_ACK" def __init__( self, @@ -207,6 +210,10 @@ def __init__( self.session_failures = defaultdict(int) self.failed_sessions = set() self.session_lock = threading.Lock() + self.room_abort_targets: dict[int, tuple[str, int]] = {} + self.room_abort_ack_sent: set[int] = set() + self.room_active_transfers: dict[int, int] = defaultdict(int) + self.room_abort_lock = threading.Lock() # Determine the number of threads to use for kv sender cpu_count = os.cpu_count() transfer_thread_pool_size = ( @@ -261,6 +268,65 @@ def __init__( def init_engine(self): self.engine = get_mooncake_transfer_engine() + def _prefill_unique_rank(self) -> int: + return ( + self.attn_tp_rank * (self.pp_size * self.attn_cp_size) + + self.pp_rank * self.attn_cp_size + + self.attn_cp_rank + ) + + def _mark_room_aborted( + self, room: int, endpoint: str, dst_port: int + ) -> None: + with self.room_abort_lock: + self.room_abort_targets[room] = (endpoint, dst_port) + + def _is_room_aborted(self, room: int) -> bool: + with self.room_abort_lock: + return room in self.room_abort_targets + + def _inc_room_active_transfer(self, room: int) -> None: + with self.room_abort_lock: + self.room_active_transfers[room] += 1 + + def _dec_room_active_transfer(self, room: int) -> None: + with self.room_abort_lock: + active = self.room_active_transfers.get(room, 0) - 1 + if active <= 0: + self.room_active_transfers.pop(room, None) + else: + self.room_active_transfers[room] = active + + def _maybe_send_abort_ack(self, room: int) -> None: + with self.room_abort_lock: + target = self.room_abort_targets.get(room) + if target is None: + return + if room in self.room_abort_ack_sent: + return + if self.room_active_transfers.get(room, 0) > 0: + return + self.room_abort_ack_sent.add(room) + + endpoint, dst_port = target + na = NetworkAddress(endpoint, dst_port) + self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart( + [ + self.ABORT_ACK_HEADER, + str(room).encode("ascii"), + str(self._prefill_unique_rank()).encode("ascii"), + ] + ) + logger.info( + "[PD_ABORT] prefill sent ABORT_ACK room=%s endpoint=%s dst_port=%s rank=%s", + room, + endpoint, + dst_port, + self._prefill_unique_rank(), + ) + self.update_status(room, KVPoll.Failed) + self.transfer_infos.pop(room, None) + def register_buffer_to_engine(self): # Batch register KV data buffers if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens: @@ -606,6 +672,12 @@ def _send_kvcache_generic( src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = ( self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs) ) + if layers_current_pp_stage > len(dst_kv_ptrs): + logger.error( + "Prefill transfer kvcache error, layers_current_pp_stage is out of range: " + f"layers_current_pp_stage={layers_current_pp_stage}, len(dst_kv_ptrs)={len(dst_kv_ptrs)}" + ) + return -1 layers_params = [ ( src_kv_ptrs[layer_id], @@ -1174,14 +1246,14 @@ def transfer_worker( if kv_chunk.room in self.transfer_infos else [] ) + if self._is_room_aborted(kv_chunk.room): + self.transfer_infos.pop(kv_chunk.room, None) + self._maybe_send_abort_ack(kv_chunk.room) + continue polls = [] dst_ranks_infos = [] # Unique id per prefill sender so decode's response set size matches expected_response_num. - prefill_unique_rank = ( - self.attn_tp_rank * (self.pp_size * self.attn_cp_size) - + self.pp_rank * self.attn_cp_size - + self.attn_cp_rank - ) + prefill_unique_rank = self._prefill_unique_rank() # When staging transfer is not yet ready (watermark/allocation pending), # the chunk is re-enqueued and we break out of the req loop to retry later. staging_deferred = False @@ -1227,23 +1299,29 @@ def transfer_worker( self.attn_tp_size == target_rank_registration_info.dst_attn_tp_size ): - if target_rank_registration_info.enable_hisparse: - ret = self.send_kvcache_hisparse( - req.mooncake_session_id, - kv_chunk.prefill_kv_indices, - target_rank_registration_info.dst_kv_ptrs, - req.dst_kv_indices, - kv_chunk.index_slice, - executor, - ) - else: - ret = self.send_kvcache( - req.mooncake_session_id, - kv_chunk.prefill_kv_indices, - target_rank_registration_info.dst_kv_ptrs, - chunked_dst_kv_indice, - executor, - ) + self._inc_room_active_transfer(kv_chunk.room) + try: + if target_rank_registration_info.enable_hisparse: + ret = self.send_kvcache_hisparse( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_rank_registration_info.dst_kv_ptrs, + req.dst_kv_indices, + kv_chunk.index_slice, + executor, + ) + else: + ret = self.send_kvcache( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_rank_registration_info.dst_kv_ptrs, + chunked_dst_kv_indice, + executor, + ) + finally: + self._dec_room_active_transfer(kv_chunk.room) + if self._is_room_aborted(kv_chunk.room): + self._maybe_send_abort_ack(kv_chunk.room) elif ( self.enable_staging and staging_strategy is not None @@ -1264,16 +1342,24 @@ def transfer_worker( # Chunk re-enqueued; stop processing remaining reqs for this chunk break else: - ret = self.send_kvcache_slice( - req.mooncake_session_id, - kv_chunk.prefill_kv_indices, - target_rank_registration_info.dst_kv_ptrs, - chunked_dst_kv_indice, - target_rank_registration_info.dst_tp_rank, - target_rank_registration_info.dst_attn_tp_size, - target_rank_registration_info.dst_kv_item_len, - executor, - ) + self._inc_room_active_transfer(kv_chunk.room) + try: + ret = self.send_kvcache_slice( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_rank_registration_info.dst_kv_ptrs, + chunked_dst_kv_indice, + target_rank_registration_info.dst_tp_rank, + target_rank_registration_info.dst_attn_tp_size, + target_rank_registration_info.dst_kv_item_len, + executor, + ) + finally: + self._dec_room_active_transfer(kv_chunk.room) + if self._is_room_aborted(kv_chunk.room): + self._maybe_send_abort_ack(kv_chunk.room) + if self._is_room_aborted(kv_chunk.room): + continue if ret != 0: with self.session_lock: self.session_failures[req.mooncake_session_id] += 1 @@ -1340,6 +1426,11 @@ def transfer_worker( if staging_deferred: continue + if self._is_room_aborted(kv_chunk.room): + self.transfer_infos.pop(kv_chunk.room, None) + self._maybe_send_abort_ack(kv_chunk.room) + continue + if ( kv_chunk.room not in self.request_status or self.check_status(kv_chunk.room) == KVPoll.Success @@ -1360,6 +1451,20 @@ def bootstrap_thread(): # KVPoll.Bootstrapping -> KVPoll.WaitingForInput while True: waiting_req_bytes = self.server_socket.recv_multipart() + if waiting_req_bytes[0] == self.ABORT_REQ_HEADER: + room = int(waiting_req_bytes[1].decode("ascii")) + endpoint = waiting_req_bytes[2].decode("ascii") + dst_port = int(waiting_req_bytes[3].decode("ascii")) + logger.info( + "[PD_ABORT] prefill recv ABORT_REQ room=%s endpoint=%s dst_port=%s", + room, + endpoint, + dst_port, + ) + self._mark_room_aborted(room, endpoint, dst_port) + self.update_status(room, KVPoll.Failed) + self._maybe_send_abort_ack(room) + continue room = waiting_req_bytes[0].decode("ascii") # Staging: decode reports consumption watermark back to prefill if room == "WATERMARK": @@ -1422,6 +1527,9 @@ def bootstrap_thread(): else: required_dst_info_num = int(waiting_req_bytes[7].decode("ascii")) room = int(room) + if self._is_room_aborted(room): + self._maybe_send_abort_ack(room) + continue if room not in self.transfer_infos: self.transfer_infos[room] = {} @@ -1485,6 +1593,27 @@ def decode_thread(): if msg[0] == b"STAGING_REQ": self._handle_staging_req(msg) continue + if msg[0] == self.ABORT_ACK_HEADER: + bootstrap_room = int(msg[1].decode("ascii")) + prefill_rank = int(msg[2].decode("ascii")) + logger.info( + "[PD_ABORT] decode recv ABORT_ACK room=%s prefill_rank=%s", + bootstrap_room, + prefill_rank, + ) + self.prefill_abort_tracker[bootstrap_room].add(prefill_rank) + expected_response_num = self.required_prefill_response_num_table.get( + bootstrap_room, 1 + ) + if ( + len(self.prefill_abort_tracker[bootstrap_room]) + >= expected_response_num + ): + self.record_failure( + bootstrap_room, "Aborted by decode request." + ) + self.update_status(bootstrap_room, KVPoll.Failed) + continue bootstrap_room, status, prefill_rank = msg status = int(status.decode("ascii")) @@ -1772,6 +1901,7 @@ def __init__( ): self.session_id = mgr.get_session_id() self.init_time = None + self.abort_requested = False super().__init__(mgr, bootstrap_addr, bootstrap_room) def _register_kv_args(self): @@ -1930,6 +2060,9 @@ def clear(self) -> None: if self.bootstrap_room in self.kv_mgr.prefill_response_tracker: self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room) + if self.bootstrap_room in self.kv_mgr.prefill_abort_tracker: + self.kv_mgr.prefill_abort_tracker.pop(self.bootstrap_room) + def failure_exception(self): # Explicitly set the status to failure since this request has failed in another rank if self.conclude_state is None: @@ -1944,12 +2077,53 @@ def failure_exception(self): raise KVTransferError(self.bootstrap_room, failure_reason) def abort(self): - self.kv_mgr.record_failure( + if self.abort_requested: + return + self.abort_requested = True + self.kv_mgr.record_failure(self.bootstrap_room, "Aborted by AbortReq.") + logger.info( + "[PD_ABORT] decode notify prefill abort start room=%s bootstrap_addr=%s", self.bootstrap_room, - "Aborted by AbortReq.", + self.bootstrap_addr, ) - # Explicitly set the status to failure since this request has been aborted - self.conclude_state = KVPoll.Failed + if not hasattr(self, "prefill_info"): + logger.error( + "MooncakeKVReceiver.abort called before init for room=%s", + self.bootstrap_room, + ) + return + try: + response = requests.post( + f"http://{self.bootstrap_addr}/abort_room", + json={ + "bootstrap_room": self.bootstrap_room, + "endpoint": self.kv_mgr.local_ip, + "dst_port": self.kv_mgr.rank_port, + "prefill_dp_rank": self.prefill_dp_rank, + "target_tp_ranks": self.target_tp_ranks, + "target_cp_ranks": self.target_cp_ranks, + "target_pp_ranks": self.target_pp_ranks, + }, + timeout=5, + ) + if response.status_code != 200: + logger.error( + "Failed to notify prefill abort for room=%s: %s %s", + self.bootstrap_room, + response.status_code, + response.text, + ) + else: + logger.info( + "[PD_ABORT] decode notify prefill abort ok room=%s", + self.bootstrap_room, + ) + except Exception as e: + logger.error( + "Failed to notify prefill abort for room=%s: %s", + self.bootstrap_room, + e, + ) class MooncakeKVBootstrapServer(CommonKVBootstrapServer): diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1619e3e28292..eaff92a1dffa 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -355,6 +355,74 @@ def maybe_prefetch_staging_for_batch(self: Scheduler, batch: ScheduleBatch) -> N if room is not None and room in kv_mgr.transfer_infos: prefetch(room) + def abort_disagg_prefill_req_if_needed(self: Scheduler, req: Req) -> bool: + kv_mgr = self.disagg_prefill_bootstrap_queue.kv_manager + is_room_aborted = getattr(kv_mgr, "_is_room_aborted", None) + if is_room_aborted is None or not is_room_aborted(req.bootstrap_room): + return False + + if getattr(req, "_disagg_prefill_abort_processed", False): + return True + + req._disagg_prefill_abort_processed = True + error_message = ( + f"Prefill request aborted by decode for {req.rid=} " + f"{req.bootstrap_room=}" + ) + logger.warning(error_message) + + if hasattr(req, "disagg_kv_sender") and hasattr(req.disagg_kv_sender, "abort"): + req.disagg_kv_sender.abort() + + req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) + if req.req_pool_idx is not None: + release_kv_cache(req, self.tree_cache) + release_req_to_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator) + prepare_abort(req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR) + req.time_stats.set_completion_time() + self.stream_output([req], req.return_logprob, None) + + if self.chunked_req is not None and self.chunked_req.rid == req.rid: + self.chunked_req = None + if self.enable_hicache_storage: + self.tree_cache.release_aborted_request(req.rid) + + return True + + def process_aborted_disagg_prefill_reqs(self: Scheduler) -> None: + aborted_reqs: List[Req] = [] + + def collect(req: Optional[Req]) -> bool: + if req is None: + return False + if self.abort_disagg_prefill_req_if_needed(req): + aborted_reqs.append(req) + return True + return False + + self.disagg_prefill_bootstrap_queue.queue = [ + req for req in self.disagg_prefill_bootstrap_queue.queue if not collect(req) + ] + self.waiting_queue = [req for req in self.waiting_queue if not collect(req)] + self.disagg_prefill_inflight_queue = [ + req for req in self.disagg_prefill_inflight_queue if not collect(req) + ] + + collect(self.chunked_req) + + batches = [self.cur_batch, self.running_batch, self.last_batch] + for attr in ("mbs", "running_mbs", "last_mbs"): + batches.extend(getattr(self, attr, []) or []) + + for batch in batches: + if batch is None or batch.is_empty(): + continue + before = batch.batch_size() + for req in list(batch.reqs): + collect(req) + if aborted_reqs and batch.batch_size() == before: + batch.filter_batch(chunked_req_to_exclude=aborted_reqs) + def get_next_disagg_prefill_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: @@ -384,6 +452,7 @@ def event_loop_normal_disagg_prefill(self: Scheduler) -> None: self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) + self.process_aborted_disagg_prefill_reqs() if self._engine_paused: continue @@ -417,6 +486,7 @@ def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) + self.process_aborted_disagg_prefill_reqs() if self._engine_paused: continue @@ -498,6 +568,9 @@ def process_batch_result_disagg_prefill( for i, (req, next_token_id) in enumerate( zip(batch.reqs, next_token_ids, strict=True) ): + if self.abort_disagg_prefill_req_if_needed(req): + continue + if req.is_chunked <= 0: req.time_stats.set_prefill_finished_time() @@ -713,17 +786,22 @@ def get_transferred_rids(self: Scheduler) -> List[str]: def process_prefill_chunk(self: Scheduler) -> None: chunked_req_to_exclude = set() if self.chunked_req: - chunked_req_to_exclude.add(self.chunked_req) - maybe_cache_unfinished_req(self.chunked_req, self.tree_cache, chunked=True) - if self.enable_overlap: - # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved - self.chunked_req.tmp_end_idx = min( - len(self.chunked_req.fill_ids), - len(self.chunked_req.origin_input_ids), - ) + if self.abort_disagg_prefill_req_if_needed(self.chunked_req): + self.running_batch.batch_is_full = False else: - self.send_kv_chunk(self.chunked_req) - self.running_batch.batch_is_full = False + chunked_req_to_exclude.add(self.chunked_req) + maybe_cache_unfinished_req( + self.chunked_req, self.tree_cache, chunked=True + ) + if self.enable_overlap: + # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved + self.chunked_req.tmp_end_idx = min( + len(self.chunked_req.fill_ids), + len(self.chunked_req.origin_input_ids), + ) + else: + self.send_kv_chunk(self.chunked_req) + self.running_batch.batch_is_full = False if self.last_batch and self.last_batch.forward_mode.is_extend(): if self.last_batch.chunked_req: @@ -747,6 +825,16 @@ def send_kv_chunk( """ Send a prefilled chunk to the decode server """ + kv_mgr = self.disagg_prefill_bootstrap_queue.kv_manager + is_room_aborted = getattr(kv_mgr, "_is_room_aborted", None) + if is_room_aborted is not None and is_room_aborted(req.bootstrap_room): + logger.info( + "Skip sending kv chunk for aborted request %s room=%s", + req.rid, + req.bootstrap_room, + ) + return + page_size = self.token_to_kv_pool_allocator.page_size start_idx = req.start_send_idx end_idx = ( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 22710db85e21..e0cfe50ab428 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3563,12 +3563,14 @@ def abort_request(self, recv_req: AbortReq): for decode_req in self.disagg_decode_prealloc_queue.queue: if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") + decode_req.req.to_finish = FINISH_ABORT() decode_req.kv_receiver.abort() # Abort requests waiting for kvcache to release tree cache for decode_req in self.disagg_decode_transfer_queue.queue: if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") + decode_req.req.to_finish = FINISH_ABORT() decode_req.kv_receiver.abort() # Abort requests already retracted to CPU cache diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 939c83f3c2f6..2daab4a455ab 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -216,6 +216,7 @@ def event_loop_pp_disagg_prefill(self: Scheduler): recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) + self.process_aborted_disagg_prefill_reqs() if not self.pp_group.is_last_rank: self._pp_commit_comm_work(self.send_req_work)