diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index 90b0208464d4..a7e67f63d9c8 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/python/sglang/srt/disaggregation/ascend/conn.py @@ -12,6 +12,7 @@ MooncakeKVManager, MooncakeKVReceiver, MooncakeKVSender, + _submit_transfer_futures, ) from sglang.srt.utils.network import get_local_ip_auto @@ -156,21 +157,11 @@ def process_layers(layers_params: List[Tuple[int, int, int]]) -> int: return self._transfer_data(mooncake_session_id, transfer_blocks) if self.enable_custom_mem_pool: - futures = [ - executor.submit( - process_layer, - src_ptr, - dst_ptr, - item_len, - ) + calls = [ + (process_layer, (src_ptr, dst_ptr, item_len)) for (src_ptr, dst_ptr, item_len) in layers_params ] - for future in concurrent.futures.as_completed(futures): - status = future.result() - if status != 0: - for f in futures: - f.cancel() - return status + return _submit_transfer_futures(executor, calls) else: # Combining all layers' params in one batch transfer is more efficient # compared to using multiple threads diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 1cbe4e7f7e08..015b63fbd668 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -156,6 +156,15 @@ def failure_exception(self): """ ... + def begin_failure_quiescence(self): + pass + + def is_transfer_quiesced(self) -> bool: + return True + + def is_failure_quiescing(self) -> bool: + return False + class BaseKVReceiver(ABC): @abstractmethod @@ -215,6 +224,15 @@ def abort(self): """ pass + def begin_failure_quiescence(self): + pass + + def is_transfer_quiesced(self) -> bool: + return True + + def is_failure_quiescing(self) -> bool: + return False + class BaseKVBootstrapServer(ABC): @abstractmethod diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 9874516065de..375c81c8e511 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -75,6 +75,7 @@ class PrefillServerInfo: page_size: Optional[int] kv_cache_dtype: Optional[str] follow_bootstrap_room: bool + capabilities: Optional[List[str]] = None # Pre-computed rank mapping (set by try_ensure_parallel_info on decode side) target_tp_rank: Optional[int] = None @@ -94,16 +95,19 @@ def __post_init__(self): str(self.kv_cache_dtype) if self.kv_cache_dtype is not None else None ) self.follow_bootstrap_room = bool(self.follow_bootstrap_room) + self.capabilities = list(self.capabilities or []) @dataclasses.dataclass class PrefillRankInfo: rank_ip: str rank_port: int + capabilities: Optional[List[str]] = None def __post_init__(self): self.rank_ip = str(self.rank_ip) self.rank_port = int(self.rank_port) + self.capabilities = list(self.capabilities or []) class CommonKVManager(BaseKVManager): @@ -421,6 +425,9 @@ def register_to_bootstrap(self): "kv_cache_dtype": self.server_args.kv_cache_dtype, "load_balance_method": self.server_args.load_balance_method, } + capabilities = self.get_bootstrap_capabilities() + if capabilities: + payload["capabilities"] = capabilities max_retries, initial_delay, max_delay = 5, 1.0, 30.0 for attempt in range(max_retries): @@ -450,6 +457,9 @@ def register_to_bootstrap(self): f"Prefill instance failed to register to bootstrap server after {max_retries} retries" ) + def get_bootstrap_capabilities(self) -> List[str]: + return [] + def _connect(self, endpoint: str, is_ipv6: bool = False): with self._socket_lock: sock = self._socket_cache.get(endpoint) @@ -1230,6 +1240,7 @@ def __init__(self, host: str, port: int): self.page_size = None self.kv_cache_dtype: Optional[str] = None self.follow_bootstrap_room: Optional[bool] = None + self.capabilities: List[str] = [] self.prefill_port_table: Dict[ int, Dict[int, Dict[int, Dict[int, PrefillRankInfo]]] ] = {} @@ -1296,6 +1307,7 @@ async def _handle_route_put(self, request: web.Request): rank_port = int(data["rank_port"]) page_size = int(data["page_size"]) kv_cache_dtype = data["kv_cache_dtype"] + capabilities = list(data.get("capabilities") or []) if self.attn_tp_size is None: self.attn_tp_size = attn_tp_size @@ -1315,6 +1327,9 @@ async def _handle_route_put(self, request: web.Request): if self.kv_cache_dtype is None and kv_cache_dtype is not None: self.kv_cache_dtype = kv_cache_dtype + if capabilities: + self.capabilities = capabilities + if self.follow_bootstrap_room is None: load_balance_method = data.get( "load_balance_method", "follow_bootstrap_room" @@ -1335,6 +1350,7 @@ async def _handle_route_put(self, request: web.Request): tp_group_table[pp_rank] = PrefillRankInfo( rank_ip=rank_ip, rank_port=rank_port, + capabilities=capabilities, ) self._registered_count += 1 @@ -1384,8 +1400,12 @@ async def _handle_route_get(self, request: web.Request): if self.follow_bootstrap_room is not None else True ), + capabilities=self.capabilities, ) - return web.json_response(dataclasses.asdict(info), status=200) + response = dataclasses.asdict(info) + if not response["capabilities"]: + response.pop("capabilities") + return web.json_response(response, status=200) if not self._is_ready(): return web.Response( @@ -1407,7 +1427,10 @@ async def _handle_route_get(self, request: web.Request): status=404, ) - return web.json_response(dataclasses.asdict(bootstrap_info), status=200) + response = dataclasses.asdict(bootstrap_info) + if not response["capabilities"]: + response.pop("capabilities") + return web.json_response(response, status=200) async def _handle_register_dp_rank(self, request: web.Request): data = await request.json() diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index a0b6bcbfc9d0..9a14bc844962 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -12,6 +12,7 @@ import logging import struct import threading +from collections import defaultdict from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -80,14 +81,20 @@ def __init__( self.scheduler = scheduler self._room_to_decode_req: dict = {} self._wm_subscribers: dict = {} + self._aborting_rooms: set = set() + self._abort_finalized_rooms: set = set() + self._abort_finalizing_rooms: set = set() + self._scatter_submitting: dict = defaultdict(int) + self._abort_lock = threading.RLock() def register_wm_subscriber(self, receiver, session_id: str) -> None: """Register a prefill's bootstrap connection for watermark broadcasts.""" if receiver is None or not getattr(receiver, "bootstrap_infos", None): return key = tuple(str(bi) for bi in receiver.bootstrap_infos) - if key not in self._wm_subscribers: - self._wm_subscribers[key] = (receiver, session_id) + with self._abort_lock: + if key not in self._wm_subscribers: + self._wm_subscribers[key] = (receiver, session_id) def num_writers_for(self, decode_req) -> int: """Compute num_writers for a specific request based on its prefill TP.""" @@ -132,11 +139,175 @@ def create(cls, kv_manager, scheduler, tp_rank: int) -> DecodeStagingHandler: # Registration: called from main thread (DecodeTransferQueue) # ------------------------------------------------------------------ - def register_decode_req(self, room: int, decode_req: DecodeRequest) -> None: - self._room_to_decode_req[room] = decode_req + def register_room_bootstrap(self, room, bootstrap_infos, receiver) -> bool: + with self._abort_lock: + if room in self._aborting_rooms: + return False + self.kv_manager._staging_ctx.room_bootstrap[room] = bootstrap_infos + self.kv_manager._staging_ctx.room_receivers[room] = receiver + return True + + def register_decode_req(self, room: int, decode_req: DecodeRequest) -> bool: + with self._abort_lock: + self._room_to_decode_req[room] = decode_req + return room not in self._aborting_rooms def unregister_decode_req(self, room: int) -> None: - self._room_to_decode_req.pop(room, None) + with self._abort_lock: + self._room_to_decode_req.pop(room, None) + self.kv_manager._staging_ctx.room_bootstrap.pop(room, None) + self.kv_manager._staging_ctx.room_receivers.pop(room, None) + writer_counts = getattr(self.kv_manager, "_chunk_writer_counts", None) + if writer_counts is not None: + writer_counts.pop(room, None) + self._scatter_submitting.pop(room, None) + + def begin_abort(self, room: int) -> None: + with self._abort_lock: + self._aborting_rooms.add(room) + + def finalize_abort(self, room: int, decode_req: DecodeRequest) -> bool: + with self._abort_lock: + if room in self._abort_finalized_rooms: + return True + if room in self._abort_finalizing_rooms or self._scatter_submitting[room]: + return False + + self.advance_scatter(decode_req) + + with self._abort_lock: + if room in self._abort_finalized_rooms: + return True + if self._scatter_submitting[room]: + return False + if getattr(decode_req, "_chunk_events", None): + return False + if getattr(decode_req, "_scatter_event", None) is not None: + return False + if getattr(decode_req, "_freeing_staging_alloc_ids", None): + return False + + receiver = decode_req.kv_receiver + alloc_ids = [ + alloc_id + for alloc_id, _offset, _round, _end, _pages in ( + getattr(receiver, "chunk_staging_infos", None) or [] + ) + if alloc_id >= 0 + ] + self._abort_finalizing_rooms.add(room) + + try: + for alloc_id in alloc_ids: + self._free_and_send_watermark(alloc_id, decode_req) + except Exception: + with self._abort_lock: + self._abort_finalizing_rooms.discard(room) + raise + + with self._abort_lock: + receiver.chunk_staging_infos = [] + writer_counts = getattr(self.kv_manager, "_chunk_writer_counts", None) + if writer_counts is not None: + writer_counts.pop(room, None) + self.kv_manager._staging_ctx.room_bootstrap.pop(room, None) + self.kv_manager._staging_ctx.room_receivers.pop(room, None) + self._room_to_decode_req.pop(room, None) + self._scatter_submitting.pop(room, None) + self._abort_finalizing_rooms.discard(room) + self._abort_finalized_rooms.add(room) + return True + + def finalize_room_abort(self, room: int) -> bool: + with self._abort_lock: + if room in self._abort_finalized_rooms: + return True + decode_req = self._room_to_decode_req.get(room) + return decode_req is not None and self.finalize_abort(room, decode_req) + + def clear_abort(self, room: int) -> None: + with self._abort_lock: + self._abort_finalized_rooms.discard(room) + + def is_abort_finalized(self, room: int) -> bool: + with self._abort_lock: + return room in self._abort_finalized_rooms + + def handle_staging_req( + self, + msg, + kv_args, + attn_tp_size: int, + kv_buffer_tensors, + ) -> bool: + room = int(msg[1].decode("ascii")) + chunk_idx = int(msg[2].decode("ascii")) + session_id = msg[4].decode("ascii") + + with self._abort_lock: + if room in self._aborting_rooms: + return False + receiver = self.kv_manager._staging_ctx.room_receivers.get(room) + decode_req = self._room_to_decode_req.get(room) + if decode_req is None or receiver is None: + return False + if decode_req.kv_receiver is not receiver: + return False + prefill_attn_tp_size = receiver.prefill_info.attn_tp_size + infos = getattr(receiver, "chunk_staging_infos", []) + existing = _get_staging_response(infos, chunk_idx) + + candidate = None + if existing is None: + candidate = _allocate_staging_request( + msg, + self.staging_allocator, + kv_args, + attn_tp_size, + prefill_attn_tp_size, + ) + + raced_alloc_id = None + with self._abort_lock: + if ( + room in self._aborting_rooms + or self.kv_manager._staging_ctx.room_receivers.get(room) is not receiver + or self._room_to_decode_req.get(room) is not decode_req + or decode_req.kv_receiver is not receiver + ): + if candidate is not None and candidate[0] >= 0: + raced_alloc_id = candidate[0] + response = None + else: + existing = _get_staging_response(infos, chunk_idx) + if existing is not None: + if candidate is not None and candidate[0] >= 0: + raced_alloc_id = candidate[0] + response = existing + else: + _set_staging_response(infos, chunk_idx, candidate) + response = candidate[1:4] + bootstrap_infos = list( + self.kv_manager._staging_ctx.room_bootstrap.get(room, ()) + ) + + if raced_alloc_id is not None: + self.staging_allocator.free(raced_alloc_id) + if response is None: + return False + + with self._abort_lock: + if room in self._aborting_rooms: + return False + _send_staging_response( + receiver, + bootstrap_infos, + room, + chunk_idx, + session_id, + response, + ) + return True # ------------------------------------------------------------------ # Scatter submission: called from decode_thread (background) @@ -150,32 +321,37 @@ def submit_chunk_scatter( Called from decode_thread. Records a CUDA event on decode_req so the main thread can later check completion and free the allocation. """ - decode_req = self._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "[STAGING] submit_chunk_scatter: room=%s not registered, " - "chunk_idx=%s. This should not happen if register_decode_req " - "is called at kv_receiver.init() time.", - room, - chunk_idx, - ) - return False - chunk_infos = getattr(decode_req.kv_receiver, "chunk_staging_infos", []) - if chunk_idx >= len(chunk_infos): - return False - alloc_id, staging_offset, _, _, _ = chunk_infos[chunk_idx] - if staging_offset < 0 or alloc_id < 0: - return False - - ok = self._scatter_region(staging_offset, page_start, num_pages, decode_req) - if ok: - event = torch.cuda.Event() - event.record(self.staging_allocator._scatter_stream) - if not hasattr(decode_req, "_chunk_events"): - decode_req._chunk_events = [] - decode_req._chunk_events.append((event, alloc_id)) - chunk_infos[chunk_idx] = (-1, -1, 0, -1, 0) - else: + with self._abort_lock: + if room in self._aborting_rooms: + return False + decode_req = self._room_to_decode_req.get(room) + if decode_req is None: + return False + chunk_infos = getattr(decode_req.kv_receiver, "chunk_staging_infos", []) + if chunk_idx >= len(chunk_infos): + return False + alloc_id, staging_offset, _, _, _ = chunk_infos[chunk_idx] + if staging_offset < 0 or alloc_id < 0: + return False + self._scatter_submitting[room] += 1 + + ok = False + event = None + try: + ok = self._scatter_region(staging_offset, page_start, num_pages, decode_req) + if ok: + event = torch.cuda.Event() + event.record(self.staging_allocator._scatter_stream) + finally: + with self._abort_lock: + self._scatter_submitting[room] -= 1 + if event is not None: + if not hasattr(decode_req, "_chunk_events"): + decode_req._chunk_events = [] + decode_req._chunk_events.append((event, alloc_id)) + chunk_infos[chunk_idx] = (-1, -1, 0, -1, 0) + + if not ok: logger.warning( "submit_chunk_scatter failed room=%s chunk_idx=%s tp_rank=%s", room, @@ -186,7 +362,8 @@ def submit_chunk_scatter( def is_staging_room(self, room: int) -> bool: """Check if a room is registered for staging scatter.""" - return room in self._room_to_decode_req + with self._abort_lock: + return room in self._room_to_decode_req def handle_chunk_arrived( self, @@ -203,21 +380,41 @@ def handle_chunk_arrived( once all writers for this chunk have reported in. Returns True if scatter was submitted. """ - chunk_writer_counts[room][chunk_idx].append((page_start, num_pages, writer_id)) - decode_req = self._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "Staging chunk arrived for unregistered room=%s chunk=%d, skipping", - room, - chunk_idx, - ) - return False - writers_arrived = len(chunk_writer_counts[room][chunk_idx]) - num_writers = self.num_writers_for(decode_req) - if writers_arrived >= num_writers: - self.submit_chunk_scatter(room, chunk_idx, page_start, num_pages) - del chunk_writer_counts[room][chunk_idx] - return True + with self._abort_lock: + if room in self._aborting_rooms: + return False + decode_req = self._room_to_decode_req.get(room) + if decode_req is None: + return False + arrivals = chunk_writer_counts[room][chunk_idx] + for existing_start, existing_pages, existing_writer in arrivals: + if existing_writer == writer_id: + if (existing_start, existing_pages) != (page_start, num_pages): + logger.error( + "Conflicting CHUNK_READY from writer=%s room=%s chunk=%s", + writer_id, + room, + chunk_idx, + ) + return False + if (existing_start, existing_pages) != (page_start, num_pages): + logger.error( + "Conflicting CHUNK_READY range room=%s chunk=%s: " + "expected=(%s,%s) got=(%s,%s)", + room, + chunk_idx, + existing_start, + existing_pages, + page_start, + num_pages, + ) + return False + arrivals.append((page_start, num_pages, writer_id)) + ready = len(arrivals) >= self.num_writers_for(decode_req) + if ready: + chunk_writer_counts[room].pop(chunk_idx, None) + if ready: + return self.submit_chunk_scatter(room, chunk_idx, page_start, num_pages) return False def submit_last_scatter_async(self, room: int) -> bool: @@ -227,24 +424,30 @@ def submit_last_scatter_async(self, room: int) -> bool: ``_staging_last_scatter_submitted`` so the main thread sees the event when it checks the flag (CPython GIL guarantees ordering). """ - decode_req = self._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "[STAGING] submit_last_scatter_async: room=%s not registered. " - "This should not happen if register_decode_req is called at " - "kv_receiver.init() time.", - room, - ) - return False - alloc_id = self._submit_last_scatter(decode_req) - if alloc_id >= 0: - event = torch.cuda.Event() - event.record(self.staging_allocator._scatter_stream) - decode_req._scatter_event = event - decode_req._scatter_alloc_id = alloc_id - decode_req._staging_last_scatter_submitted = True - else: - decode_req._staging_scatter_done = True + with self._abort_lock: + if room in self._aborting_rooms: + return False + decode_req = self._room_to_decode_req.get(room) + if decode_req is None: + return False + self._scatter_submitting[room] += 1 + + alloc_id = -1 + event = None + try: + alloc_id = self._submit_last_scatter(decode_req) + if alloc_id >= 0: + event = torch.cuda.Event() + event.record(self.staging_allocator._scatter_stream) + finally: + with self._abort_lock: + self._scatter_submitting[room] -= 1 + if event is not None: + decode_req._scatter_event = event + decode_req._scatter_alloc_id = alloc_id + decode_req._staging_last_scatter_submitted = True + elif alloc_id < 0: + decode_req._staging_scatter_done = True return True # ------------------------------------------------------------------ @@ -253,9 +456,10 @@ def submit_last_scatter_async(self, room: int) -> bool: def is_done(self, decode_req: DecodeRequest) -> bool: """Return True if staging scatter is complete for this request.""" - if not getattr(decode_req, "_staging_scatter_done", False): - return False - return not getattr(decode_req, "_chunk_events", None) + with self._abort_lock: + if not getattr(decode_req, "_staging_scatter_done", False): + return False + return not getattr(decode_req, "_chunk_events", None) def advance_scatter(self, decode_req: DecodeRequest) -> None: """Check CUDA events and free completed staging allocations. @@ -264,24 +468,28 @@ def advance_scatter(self, decode_req: DecodeRequest) -> None: (via submit_chunk_scatter / submit_last_scatter_async). This method only polls the recorded events and releases staging memory. """ - room = decode_req.req.bootstrap_room - chunk_events = getattr(decode_req, "_chunk_events", None) - if chunk_events: - for i in range(len(chunk_events) - 1, -1, -1): - event, alloc_id = chunk_events[i] - if event.query(): - chunk_events.pop(i) - self._free_and_send_watermark(alloc_id, decode_req) - - if not getattr(decode_req, "_staging_last_scatter_submitted", False): - return - - event = getattr(decode_req, "_scatter_event", None) - if event is not None and event.query(): - self._free_and_send_watermark(decode_req._scatter_alloc_id, decode_req) - decode_req._scatter_event = None - decode_req._scatter_alloc_id = -1 - decode_req._staging_scatter_done = True + with self._abort_lock: + chunk_events = list(getattr(decode_req, "_chunk_events", None) or []) + last_event = getattr(decode_req, "_scatter_event", None) + last_alloc_id = getattr(decode_req, "_scatter_alloc_id", -1) + + completed_chunks = [item for item in chunk_events if item[0].query()] + last_complete = last_event is not None and last_event.query() + alloc_ids = [] + with self._abort_lock: + current_events = getattr(decode_req, "_chunk_events", None) or [] + for item in completed_chunks: + if item in current_events: + current_events.remove(item) + alloc_ids.append(item[1]) + if last_complete and decode_req._scatter_event is last_event: + decode_req._scatter_event = None + decode_req._scatter_alloc_id = -1 + decode_req._staging_scatter_done = True + alloc_ids.append(last_alloc_id) + + for alloc_id in alloc_ids: + self._free_and_send_watermark(alloc_id, decode_req) # ------------------------------------------------------------------ # Internal methods @@ -373,13 +581,34 @@ def _free_and_send_watermark( self, alloc_id: int, decode_req: DecodeRequest ) -> None: """Free a staging allocation and broadcast watermark to all prefills.""" - self.staging_allocator.free(alloc_id) - post_wm = self.staging_allocator.get_watermark() - room = decode_req.req.bootstrap_room + with self._abort_lock: + freed = getattr(decode_req, "_freed_staging_alloc_ids", None) + if freed is None: + freed = decode_req._freed_staging_alloc_ids = set() + freeing = getattr(decode_req, "_freeing_staging_alloc_ids", None) + if freeing is None: + freeing = decode_req._freeing_staging_alloc_ids = set() + if alloc_id in freed or alloc_id in freeing: + return + freeing.add(alloc_id) + + try: + self.staging_allocator.free(alloc_id) + post_wm = self.staging_allocator.get_watermark() + except Exception: + with self._abort_lock: + freeing.discard(alloc_id) + raise + + with self._abort_lock: + freeing.discard(alloc_id) + freed.add(alloc_id) + subscribers = list(self._wm_subscribers.values()) + wm_round, wm_tail = post_wm wm_round_b = str(wm_round).encode("ascii") wm_tail_b = str(wm_tail).encode("ascii") - for _key, (receiver, session_id) in list(self._wm_subscribers.items()): + for receiver, session_id in subscribers: sid_b = session_id.encode("ascii") for bootstrap_info in receiver.bootstrap_infos: try: @@ -654,125 +883,134 @@ def init_staging_allocator(register_fn, kv_args): return allocator -def handle_staging_req( +def _get_staging_response(infos, chunk_idx): + from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator + + if chunk_idx < len(infos) and infos[chunk_idx][0] >= 0: + _, offset, rnd, end, _ = infos[chunk_idx] + return offset, rnd, end + if ( + chunk_idx < len(infos) + and infos[chunk_idx][1] == StagingAllocator.ALLOC_OVERSIZED + ): + return StagingAllocator.ALLOC_OVERSIZED, 0, -1 + return None + + +def _allocate_staging_request( msg, staging_allocator, kv_args, attn_tp_size: int, prefill_attn_tp_size: int, - kv_buffer_tensors, - room_receivers: dict, - room_bootstrap: dict, ): - """Allocate staging for a chunk on-demand and send STAGING_RSP to prefill. - - Deduplicates: multiple prefill TP ranks requesting the same (room, chunk_idx) - only allocate once. Sends ALLOC_OVERSIZED on permanent failure. - """ - from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator + from sglang.srt.disaggregation.common.staging_buffer import ( + StagingAllocator, + compute_staging_layout, + resolve_total_kv_heads, + ) room = int(msg[1].decode("ascii")) chunk_idx = int(msg[2].decode("ascii")) chunk_num_pages = int(msg[3].decode("ascii")) - session_id = msg[4].decode("ascii") - - if staging_allocator is None: - logger.warning( - "STAGING_REQ ignored: allocator is None room=%s chunk=%s", + page_size = kv_args.page_size + kv_item_lens = kv_args.kv_item_lens + num_kv_layers = len(kv_item_lens) // 2 + decode_bytes_per_token = kv_item_lens[0] // page_size + total_kv_heads = resolve_total_kv_heads(kv_args, attn_tp_size) + dst_heads_per_rank = max(1, total_kv_heads // max(1, attn_tp_size)) + bytes_per_head_per_token = decode_bytes_per_token // dst_heads_per_rank + dst_tp_rank = kv_args.engine_rank % max(1, attn_tp_size) + + _, _, required = compute_staging_layout( + prefill_attn_tp_size, + attn_tp_size, + dst_tp_rank, + total_kv_heads, + chunk_num_pages * page_size, + bytes_per_head_per_token, + num_kv_layers, + ) + result = staging_allocator.assign(required) + if result is None: + logger.error( + "[STAGING_REQ] alloc failed room=%s chunk=%d (need %d bytes, " + "buffer total=%d bytes). Increase SGLANG_DISAGG_STAGING_POOL_SIZE_MB.", room, chunk_idx, + required, + staging_allocator.total_size, ) - return + return -1, StagingAllocator.ALLOC_OVERSIZED, 0, -1, chunk_num_pages + alloc_id, offset, rnd = result + return alloc_id, offset, rnd, offset + required, chunk_num_pages + + +def _set_staging_response(infos, chunk_idx, allocation) -> None: + while len(infos) <= chunk_idx: + infos.append((-1, -1, 0, -1, 0)) + infos[chunk_idx] = allocation + + +def _send_staging_response( + receiver, bootstrap_infos, room, chunk_idx, session_id, response +) -> None: + offset, rnd, end = response + for bootstrap_info in bootstrap_infos: + try: + sock, lock = receiver._connect_to_bootstrap_server(bootstrap_info) + with lock: + sock.send_multipart( + [ + b"STAGING_RSP", + str(room).encode("ascii"), + str(chunk_idx).encode("ascii"), + str(offset).encode("ascii"), + str(rnd).encode("ascii"), + str(end).encode("ascii"), + session_id.encode("ascii"), + ] + ) + except Exception: + pass + +def handle_staging_req( + msg, + staging_allocator, + kv_args, + attn_tp_size: int, + prefill_attn_tp_size: int, + kv_buffer_tensors, + room_receivers: dict, + room_bootstrap: dict, +): + room = int(msg[1].decode("ascii")) + chunk_idx = int(msg[2].decode("ascii")) + session_id = msg[4].decode("ascii") receiver = room_receivers.get(room) - if receiver is None: - logger.warning( - "STAGING_REQ dropped: no receiver for room=%s chunk=%s session=%s", - room, - chunk_idx, - session_id, - ) + if staging_allocator is None or receiver is None: return infos = getattr(receiver, "chunk_staging_infos", []) - - if chunk_idx < len(infos) and infos[chunk_idx][0] >= 0: - _, offset, rnd, end, _ = infos[chunk_idx] - elif ( - chunk_idx < len(infos) - and infos[chunk_idx][1] == StagingAllocator.ALLOC_OVERSIZED - ): - offset, rnd, end = StagingAllocator.ALLOC_OVERSIZED, 0, -1 - else: - from sglang.srt.disaggregation.common.staging_buffer import ( - compute_staging_layout, - resolve_total_kv_heads, - ) - - page_size = kv_args.page_size - kv_item_lens = kv_args.kv_item_lens - num_kv_layers = len(kv_item_lens) // 2 - decode_bytes_per_token = kv_item_lens[0] // page_size - total_kv_heads = resolve_total_kv_heads(kv_args, attn_tp_size) - dst_heads_per_rank = max(1, total_kv_heads // max(1, attn_tp_size)) - bytes_per_head_per_token = decode_bytes_per_token // dst_heads_per_rank - dst_tp_rank = kv_args.engine_rank % max(1, attn_tp_size) - - chunk_tokens = chunk_num_pages * page_size - _, _, required = compute_staging_layout( - prefill_attn_tp_size, + response = _get_staging_response(infos, chunk_idx) + if response is None: + allocation = _allocate_staging_request( + msg, + staging_allocator, + kv_args, attn_tp_size, - dst_tp_rank, - total_kv_heads, - chunk_tokens, - bytes_per_head_per_token, - num_kv_layers, + prefill_attn_tp_size, ) - result = staging_allocator.assign(required) - if result is None: - logger.error( - "[STAGING_REQ] alloc failed room=%s chunk=%d (need %d bytes, " - "buffer total=%d bytes). Increase SGLANG_DISAGG_STAGING_POOL_SIZE_MB.", - room, - chunk_idx, - required, - staging_allocator.total_size, - ) - offset, rnd, end = StagingAllocator.ALLOC_OVERSIZED, 0, -1 - while len(infos) <= chunk_idx: - infos.append((-1, -1, 0, -1, 0)) - infos[chunk_idx] = ( - -1, - StagingAllocator.ALLOC_OVERSIZED, - 0, - -1, - chunk_num_pages, - ) - else: - alloc_id, offset, rnd = result - end = offset + required - while len(infos) <= chunk_idx: - infos.append((-1, -1, 0, -1, 0)) - infos[chunk_idx] = (alloc_id, offset, rnd, end, chunk_num_pages) - - bootstrap_infos = room_bootstrap.get(room) - if bootstrap_infos: - for bi in bootstrap_infos: - try: - sock, lock = receiver._connect_to_bootstrap_server(bi) - with lock: - sock.send_multipart( - [ - b"STAGING_RSP", - str(room).encode("ascii"), - str(chunk_idx).encode("ascii"), - str(offset).encode("ascii"), - str(rnd).encode("ascii"), - str(end).encode("ascii"), - session_id.encode("ascii"), - ] - ) - except Exception: - pass + _set_staging_response(infos, chunk_idx, allocation) + response = allocation[1:4] + _send_staging_response( + receiver, + room_bootstrap.get(room, ()), + room, + chunk_idx, + session_id, + response, + ) def prefetch_staging_reqs( @@ -782,6 +1020,7 @@ def prefetch_staging_reqs( chunked_prefill_size: int, staging_requested: set, prefetch_sockets: dict, + send_message=None, ) -> None: """Send STAGING_REQ for all chunks before the prefill forward starts. @@ -819,6 +1058,16 @@ def prefetch_staging_reqs( remaining = total_pages - chunk_idx * full_chunk_pages chunk_pages = min(full_chunk_pages, remaining) try: + parts = [ + b"STAGING_REQ", + str(room).encode("ascii"), + str(chunk_idx).encode("ascii"), + str(chunk_pages).encode("ascii"), + session_id.encode("ascii"), + ] + if send_message is not None: + send_message(tinfo.endpoint, tinfo.dst_port, parts) + continue na = NetworkAddress(tinfo.endpoint, tinfo.dst_port) ep = na.to_tcp() if ep not in prefetch_sockets: @@ -827,14 +1076,6 @@ def prefetch_staging_reqs( sock.setsockopt(zmq.IPV6, 1) sock.connect(ep) prefetch_sockets[ep] = sock - prefetch_sockets[ep].send_multipart( - [ - b"STAGING_REQ", - str(room).encode("ascii"), - str(chunk_idx).encode("ascii"), - str(chunk_pages).encode("ascii"), - session_id.encode("ascii"), - ] - ) + prefetch_sockets[ep].send_multipart(parts) except Exception: staging_requested.discard(stg_key) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index acf1b120b289..a5a80f3a2244 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1102,12 +1102,6 @@ def _c128_state_payload(): ) assert decode_req.metadata_buffer_index is not None page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size) - decode_req.kv_receiver.send_metadata( - page_indices, - decode_req.metadata_buffer_index, - state_indices, - decode_prefix_len=total_prefix_len, - ) if ( self.transfer_queue.enable_staging and hasattr(decode_req.kv_receiver, "require_staging") @@ -1116,6 +1110,12 @@ def _c128_state_payload(): self.transfer_queue.staging_handler.register_decode_req( decode_req.req.bootstrap_room, decode_req ) + decode_req.kv_receiver.send_metadata( + page_indices, + decode_req.metadata_buffer_index, + state_indices, + decode_prefix_len=total_prefix_len, + ) preallocated_reqs.append(decode_req) indices_to_remove.add(i) decode_req.req.time_stats.set_decode_transfer_queue_entry_time( diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 70c1df8c8193..5afe06a66078 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -164,6 +164,15 @@ def poll(self) -> KVPoll: return KVPoll.Transferring return poll + def begin_failure_quiescence(self): + return self.decode_req.kv_receiver.begin_failure_quiescence() + + def is_transfer_quiesced(self) -> bool: + return self.decode_req.kv_receiver.is_transfer_quiesced() + + def is_failure_quiescing(self) -> bool: + return self.decode_req.kv_receiver.is_failure_quiescing() + class DecodeHiCacheTransferMixin: """HiCache hooks for ``DecodeTransferQueue``: drive restore state machine.""" diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 407a58214c8e..de63d072c1e2 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -57,12 +57,96 @@ logger = logging.getLogger(__name__) +QUIESCE_CAPABILITY = "QUIESCE_V1" +QUIESCE_CAPABILITY_BYTES = QUIESCE_CAPABILITY.encode("ascii") +ABORT_RETRY_INTERVAL_S = 0.5 + FAILED_SESSION_RECOVERIES = Counter( "sglang:failed_session_recoveries_total", "Number of mooncake_session_ids un-blacklisted via probe.", ) +class RoomTransferLifetime: + def __init__(self): + self._condition = threading.Condition() + self._active = 0 + self._accepting = True + + def acquire(self) -> bool: + with self._condition: + if not self._accepting: + return False + self._active += 1 + return True + + def release(self) -> None: + with self._condition: + if self._active == 0: + raise RuntimeError("Room transfer lease released without acquire") + self._active -= 1 + if self._active == 0: + self._condition.notify_all() + + def abort(self) -> None: + with self._condition: + self._accepting = False + if self._active == 0: + self._condition.notify_all() + + def is_quiesced(self) -> bool: + with self._condition: + return not self._accepting and self._active == 0 + + def wait(self) -> None: + with self._condition: + self._condition.wait_for(lambda: self._active == 0) + + +def _wait_transfer_futures(futures) -> int: + first_status = 0 + first_exception = None + for future in concurrent.futures.as_completed(futures): + try: + status = future.result() + if status != 0 and first_status == 0: + first_status = status + for sibling in futures: + sibling.cancel() + except concurrent.futures.CancelledError: + continue + except BaseException as e: + if first_exception is None: + first_exception = e + for sibling in futures: + sibling.cancel() + concurrent.futures.wait(futures) + if first_exception is not None: + raise first_exception + return first_status + + +def _submit_transfer_futures(executor, calls) -> int: + futures = [] + try: + for fn, args in calls: + futures.append(executor.submit(fn, *args)) + except BaseException: + for future in futures: + future.cancel() + concurrent.futures.wait(futures) + raise + return _wait_transfer_futures(futures) + + +def _has_kv_registration_capability(msg: List[bytes]) -> bool: + return len(msg) > 14 and msg[14] == QUIESCE_CAPABILITY_BYTES + + +def _has_transfer_metadata_capability(msg: List[bytes]) -> bool: + return len(msg) > 9 and msg[9] == QUIESCE_CAPABILITY_BYTES + + # decode @dataclasses.dataclass class TransferInfo: @@ -76,6 +160,7 @@ class TransferInfo: required_dst_info_num: int is_dummy: bool decode_prefix_len: Optional[int] = None + abort_token: bytes = b"" # Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed) staging: Optional[StagingTransferInfo] = None @@ -104,6 +189,7 @@ def from_zmq(cls, msg: List[bytes]): decode_prefix_len=( int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None ), + abort_token=msg[10] if len(msg) > 10 else b"", ) @@ -161,15 +247,22 @@ def __init__( is_mla_backend: Optional[bool] = False, ): super().__init__(args, disaggregation_mode, server_args, is_mla_backend) + self._endpoint_send_locks = {} + self._endpoint_send_locks_lock = threading.Lock() self.init_engine() self.register_buffer_to_engine() self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.enable_trace = server_args.enable_trace if self.disaggregation_mode == DisaggregationMode.PREFILL: - self.start_prefill_thread() + self.room_lifetimes = {} + self.room_abort_tokens = defaultdict(set) + self.room_lifetimes_lock = threading.Lock() + self._abort_ack_inflight = set() + self._abort_ack_lock = threading.Lock() self.session_failures = defaultdict(int) self.failed_sessions = set() self.session_lock = threading.Lock() + self.start_prefill_thread() # Determine the number of threads to use for kv sender cpu_count = os.cpu_count() transfer_thread_pool_size = ( @@ -228,6 +321,9 @@ def __init__( daemon=True, ).start() elif self.disaggregation_mode == DisaggregationMode.DECODE: + self.aborting_rooms = set() + self.abort_receivers = {} + self._abort_receivers_lock = threading.Lock() self._staging_ctx = DecodeStagingContext() if self.enable_staging else None if self.enable_staging: self._init_staging_allocator() @@ -235,9 +331,75 @@ def __init__( self._chunk_writer_counts: dict = defaultdict(lambda: defaultdict(list)) self.start_decode_thread() + def _get_room_lifetime(self, room: int) -> RoomTransferLifetime: + with self.room_lifetimes_lock: + return self.room_lifetimes.setdefault(room, RoomTransferLifetime()) + + def _find_room_lifetime(self, room: int) -> Optional[RoomTransferLifetime]: + with self.room_lifetimes_lock: + return self.room_lifetimes.get(room) + + def _abort_room(self, room: int) -> Optional[RoomTransferLifetime]: + lifetime = self._find_room_lifetime(room) + if lifetime is not None: + lifetime.abort() + return lifetime + + def _register_room_abort_token(self, room: int, token: bytes) -> None: + with self.room_lifetimes_lock: + self.room_abort_tokens[room].add(token) + + def _abort_room_for_token( + self, room: int, token: bytes + ) -> Optional[RoomTransferLifetime]: + with self.room_lifetimes_lock: + if token not in self.room_abort_tokens.get(room, ()): + return None + lifetime = self.room_lifetimes.get(room) + if lifetime is not None: + lifetime.abort() + return lifetime + + def _clear_room_lifetime(self, room: int) -> None: + with self.room_lifetimes_lock: + lifetime = self.room_lifetimes.get(room) + if lifetime is not None and not lifetime.is_quiesced(): + return + self.room_lifetimes.pop(room, None) + self.room_abort_tokens.pop(room, None) + + def _get_endpoint_send_lock(self, endpoint: str) -> threading.Lock: + with self._endpoint_send_locks_lock: + return self._endpoint_send_locks.setdefault(endpoint, threading.Lock()) + + def _send_manager_message( + self, remote: str, dst_port: int, parts: List[bytes] + ) -> None: + na = NetworkAddress(remote, dst_port) + endpoint = na.to_tcp() + with self._get_endpoint_send_lock(endpoint): + self._connect(endpoint, is_ipv6=na.is_ipv6).send_multipart(parts) + def init_engine(self): self.engine = get_mooncake_transfer_engine() + def get_bootstrap_capabilities(self) -> List[str]: + return [QUIESCE_CAPABILITY] + + def try_ensure_parallel_info(self, bootstrap_addr: str) -> bool: + if not super().try_ensure_parallel_info(bootstrap_addr): + return False + info = self.prefill_info_table[bootstrap_addr] + if QUIESCE_CAPABILITY in info.capabilities: + return True + self.prefill_info_table.pop(bootstrap_addr, None) + logger.error( + "Mooncake prefill %s does not advertise %s; refusing KV transfer", + bootstrap_addr, + QUIESCE_CAPABILITY, + ) + return False + 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: @@ -277,8 +439,12 @@ def deregister_buffer_to_engine(self): # ------------------------------------------------------------------ def register_staging_room_bootstrap(self, room, bootstrap_infos, receiver): + handler = getattr(self, "_staging_handler", None) + if handler is not None: + return handler.register_room_bootstrap(room, bootstrap_infos, receiver) self._staging_ctx.room_bootstrap[room] = bootstrap_infos self._staging_ctx.room_receivers[room] = receiver + return True def set_kv_buffer_tensors(self, k_buffers: list, v_buffers: list, page_size: int): self.kv_buffer_tensors = { @@ -311,39 +477,46 @@ def _init_staging_allocator(self): self.kv_buffer_tensors = None def _handle_staging_req(self, msg): - from sglang.srt.disaggregation.common.staging_handler import ( - handle_staging_req, - ) - room = int(msg[1].decode("ascii")) session_id = msg[4].decode("ascii") handler = self._staging_handler assert ( handler is not None ), "STAGING_REQ received before staging handler initialized" - decode_req = handler._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "STAGING_REQ received for unregistered room=%s, skipping", - room, - ) - return - prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size - handle_staging_req( + accepted = handler.handle_staging_req( msg, - self._staging_ctx.allocator, self.kv_args, self.attn_tp_size, - prefill_tp, getattr(self, "kv_buffer_tensors", None), - self._staging_ctx.room_receivers, - self._staging_ctx.room_bootstrap, ) receiver = self._staging_ctx.room_receivers.get(room) - if receiver is not None: + if accepted and receiver is not None: handler.register_wm_subscriber(receiver, session_id) + def _handle_chunk_ready(self, msg) -> bool: + room = int(msg[1].decode("ascii")) + if room in self.aborting_rooms: + return False + handler = self._staging_handler + assert handler is not None, "CHUNK_READY before staging handler initialized" + return handler.handle_chunk_arrived( + room, + int(msg[2].decode("ascii")), + int(msg[3].decode("ascii")), + int(msg[4].decode("ascii")), + msg[6].decode("ascii") if len(msg) > 6 else msg[5].decode("ascii"), + self._chunk_writer_counts, + ) + + def _handle_abort_ack(self, msg) -> None: + room = int(msg[1].decode("ascii")) + token = msg[2] if len(msg) > 2 else None + with self._abort_receivers_lock: + receiver = self.abort_receivers.get(room) + if receiver is not None: + receiver._record_abort_ack(token) + def _is_watermark_ready( self, session_id: str, alloc_round: int, alloc_end: int ) -> bool: @@ -365,11 +538,9 @@ def _try_create_staging_strategy(self, staging_buffer): def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank): """Notify decode that a non-last staging chunk RDMA is complete.""" try: - na = NetworkAddress(req.endpoint, req.dst_port) - self._connect( - na.to_tcp(), - is_ipv6=na.is_ipv6, - ).send_multipart( + self._send_manager_message( + req.endpoint, + req.dst_port, [ b"CHUNK_READY", str(req.room).encode("ascii"), @@ -378,7 +549,7 @@ def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank): str(len(kv_chunk.prefill_kv_indices)).encode("ascii"), req.mooncake_session_id.encode("ascii"), str(prefill_unique_rank).encode("ascii"), - ] + ], ) except Exception: pass @@ -469,6 +640,7 @@ def _prefetch_staging_reqs(self, room: int): self.server_args.chunked_prefill_size, self._staging_ctx.prefetch_requested, self._staging_ctx.prefetch_sockets, + send_message=self._send_manager_message, ) def send_kvcache_staged( @@ -670,22 +842,11 @@ def process_layers(layers_params: List[Tuple[int, int, int]]) -> int: return self._transfer_data(mooncake_session_id, transfer_blocks) if self.enable_custom_mem_pool: - futures = [ - executor.submit( - process_layer, - src_ptr, - dst_ptr, - item_len, - ) + calls = [ + (process_layer, (src_ptr, dst_ptr, item_len)) for (src_ptr, dst_ptr, item_len) in layers_params ] - for future in concurrent.futures.as_completed(futures): - status = future.result() - if status != 0: - for f in futures: - f.cancel() - return status - return 0 + return _submit_transfer_futures(executor, calls) else: # Combining all layers' params in one batch transfer is more efficient # compared to using multiple threads @@ -813,24 +974,13 @@ def process_layer_tp_aware(src_layer_ptr, dst_layer_ptr): mooncake_session_id, src_addr_list, dst_addr_list, length_list ) - futures = [] + calls = [] for i in range(layers_current_pp_stage): - futures.append( - executor.submit(process_layer_tp_aware, src_k_ptrs[i], dst_k_ptrs[i]) - ) + calls.append((process_layer_tp_aware, (src_k_ptrs[i], dst_k_ptrs[i]))) for i in range(layers_current_pp_stage): - futures.append( - executor.submit(process_layer_tp_aware, src_v_ptrs[i], dst_v_ptrs[i]) - ) + calls.append((process_layer_tp_aware, (src_v_ptrs[i], dst_v_ptrs[i]))) - for future in concurrent.futures.as_completed(futures): - status = future.result() - if status != 0: - for f in futures: - f.cancel() - return status - - return 0 + return _submit_transfer_futures(executor, calls) def send_aux( self, @@ -890,10 +1040,9 @@ def send_aux_data_to_endpoint( aux_index: int, data: bytes, ): - na = NetworkAddress(remote, dst_port) - socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6) - - socket.send_multipart( + self._send_manager_message( + remote, + dst_port, [ MooncakeKVManager.AUX_DATA_HEADER, str(room).encode("ascii"), @@ -901,12 +1050,15 @@ def send_aux_data_to_endpoint( str(aux_index).encode("ascii"), struct.pack(">I", len(data)), data, - ] + ], ) def _handle_aux_data(self, msg: List[bytes]): """Handle AUX_DATA messages received by the decode thread.""" room = int(msg[1].decode("ascii")) + if room in self.aborting_rooms: + logger.debug("Dropping AUX_DATA for aborting room %s", room) + return buffer_index = int(msg[2].decode("ascii")) aux_index = int(msg[3].decode("ascii")) data_length = struct.unpack(">I", msg[4])[0] @@ -1234,13 +1386,14 @@ def _send_mamba_state_slice( def sync_status_to_decode_endpoint( self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int ): - na = NetworkAddress(remote, dst_port) - self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart( + self._send_manager_message( + remote, + dst_port, [ str(room).encode("ascii"), str(status).encode("ascii"), str(prefill_rank).encode("ascii"), - ] + ], ) def transfer_worker( @@ -1259,6 +1412,8 @@ def transfer_worker( ) while True: + lifetime = None + lease_acquired = False try: kv_chunk: TransferKVChunk = queue.get() if self.enable_trace: @@ -1268,6 +1423,12 @@ def transfer_worker( MooncakeRequestStage.MOONCAKE_WORKER_SEND.level, ) + lifetime = self._get_room_lifetime(kv_chunk.room) + lease_acquired = lifetime.acquire() + if not lease_acquired: + logger.debug("Skipping chunk for aborted room %s", kv_chunk.room) + continue + if ( kv_chunk.room not in self.request_status or self.check_status(kv_chunk.room) == KVPoll.Failed @@ -1514,6 +1675,48 @@ def transfer_worker( getattr(kv_chunk, "room", ""), ) continue + finally: + if lease_acquired: + lifetime.release() + + def _send_abort_ack_when_quiesced( + self, + lifetime: Optional[RoomTransferLifetime], + room: int, + decode_ip: str, + decode_port: int, + token: bytes, + ) -> None: + try: + if lifetime is not None: + lifetime.wait() + self._send_manager_message( + decode_ip, decode_port, [b"ABORT_ACK", str(room).encode("ascii"), token] + ) + except Exception as e: + logger.debug("Failed to send ABORT_ACK for room %s: %s", room, e) + finally: + with self._abort_ack_lock: + self._abort_ack_inflight.discard((room, decode_ip, decode_port, token)) + + def _schedule_abort_ack( + self, + lifetime: Optional[RoomTransferLifetime], + room: int, + decode_ip: str, + decode_port: int, + token: bytes, + ) -> None: + key = (room, decode_ip, decode_port, token) + with self._abort_ack_lock: + if key in self._abort_ack_inflight: + return + self._abort_ack_inflight.add(key) + threading.Thread( + target=self._send_abort_ack_when_quiesced, + args=(lifetime, room, decode_ip, decode_port, token), + daemon=True, + ).start() def start_prefill_thread(self): def bootstrap_thread(): @@ -1543,9 +1746,12 @@ 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")) + token = waiting_req_bytes[4] if len(waiting_req_bytes) > 4 else b"" + lifetime = self._abort_room_for_token(room_to_be_aborted, token) # No need to abort the room if it has already succeeded if ( - room_to_be_aborted in self.request_status + lifetime is not None + and 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) @@ -1558,26 +1764,23 @@ def bootstrap_thread(): f"Received abort notification for room {room_to_be_aborted}, " f"ignoring (already completed or unknown)" ) - # 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}" - ) + self._schedule_abort_ack( + lifetime, + room_to_be_aborted, + decode_ip, + decode_port, + token, + ) continue mooncake_session_id = waiting_req_bytes[3].decode("ascii") if room == "None": + if not _has_kv_registration_capability(waiting_req_bytes): + logger.error( + "Rejecting Mooncake KV registration without %s from %s", + QUIESCE_CAPABILITY, + mooncake_session_id, + ) + continue self.decode_kv_args_table[mooncake_session_id] = ( KVArgsRegisterInfo.from_zmq(waiting_req_bytes) ) @@ -1593,23 +1796,44 @@ def bootstrap_thread(): else: required_dst_info_num = int(waiting_req_bytes[7].decode("ascii")) room = int(room) - if room not in self.transfer_infos: - self.transfer_infos[room] = {} - - self.transfer_infos[room][mooncake_session_id] = ( - TransferInfo.from_zmq(waiting_req_bytes) - ) - # NOTE: after bootstrapping we can mark the req as waiting for input - if len(self.transfer_infos[room]) == required_dst_info_num: - self.req_to_decode_prefix_len[room] = next( - ( - info.decode_prefix_len - for info in self.transfer_infos[room].values() - if info.decode_prefix_len is not None - ), - 0, + if not _has_transfer_metadata_capability(waiting_req_bytes): + self.record_failure( + room, + f"Mooncake decode metadata is missing {QUIESCE_CAPABILITY}", + ) + self.update_status(room, KVPoll.Failed) + continue + if ( + room not in self.request_status + or self.check_status(room) == KVPoll.Success + ): + logger.debug( + "Dropping metadata for unknown/completed room %s", room ) - self.update_status(room, KVPoll.WaitingForInput) + continue + lifetime = self._get_room_lifetime(room) + if not lifetime.acquire(): + logger.debug("Dropping bootstrap for aborted room %s", room) + continue + try: + if room not in self.transfer_infos: + self.transfer_infos[room] = {} + + transfer_info = TransferInfo.from_zmq(waiting_req_bytes) + self._register_room_abort_token(room, transfer_info.abort_token) + self.transfer_infos[room][mooncake_session_id] = transfer_info + if len(self.transfer_infos[room]) == required_dst_info_num: + self.req_to_decode_prefix_len[room] = next( + ( + info.decode_prefix_len + for info in self.transfer_infos[room].values() + if info.decode_prefix_len is not None + ), + 0, + ) + self.update_status(room, KVPoll.WaitingForInput) + finally: + lifetime.release() threading.Thread(target=bootstrap_thread).start() @@ -1623,41 +1847,28 @@ def decode_thread(): # Staging: prefill notifies a chunk written to staging buffer if msg[0] == b"CHUNK_READY": - room = int(msg[1].decode("ascii")) - chunk_idx = int(msg[2].decode("ascii")) - page_start = int(msg[3].decode("ascii")) - num_pages = int(msg[4].decode("ascii")) - session_id = msg[5].decode("ascii") - handler = self._staging_handler - assert ( - handler is not None - ), "CHUNK_READY received before staging handler initialized" - handler.handle_chunk_arrived( - room, - chunk_idx, - page_start, - num_pages, - session_id, - self._chunk_writer_counts, - ) + self._handle_chunk_ready(msg) continue # Staging: prefill pre-requests staging allocation before forward if msg[0] == b"STAGING_REQ": + room = int(msg[1].decode("ascii")) + if room in self.aborting_rooms: + continue self._handle_staging_req(msg) continue # 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}") + self._handle_abort_ack(msg) continue bootstrap_room, status, prefill_rank = msg status = int(status.decode("ascii")) bootstrap_room = int(bootstrap_room.decode("ascii")) prefill_rank = int(prefill_rank.decode("ascii")) + if bootstrap_room in self.aborting_rooms: + continue if status == KVPoll.Success: if bootstrap_room in self.request_status: @@ -1698,6 +1909,33 @@ def add_transfer_request( assert self.disaggregation_mode == DisaggregationMode.PREFILL assert not is_last_chunk or (is_last_chunk and aux_index is not None) + lifetime = self._get_room_lifetime(bootstrap_room) + if not lifetime.acquire(): + logger.debug("Request with bootstrap_room=%s was aborted", bootstrap_room) + return + try: + self._add_transfer_request_with_lease( + bootstrap_room, + kv_indices, + index_slice, + is_last_chunk, + aux_index, + state_indices, + trace_ctx, + ) + finally: + lifetime.release() + + def _add_transfer_request_with_lease( + self, + bootstrap_room, + kv_indices, + index_slice, + is_last_chunk, + aux_index, + state_indices, + trace_ctx, + ): if ( bootstrap_room not in self.request_status or self.check_status(bootstrap_room) == KVPoll.Failed @@ -1795,6 +2033,7 @@ def __init__( ): super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank) self.conclude_state = None + self._failure_quiescence_started = False self.init_time = time.time() self._init_trace_ctx() @@ -1833,18 +2072,50 @@ def send( def poll(self) -> KVPoll: if self.conclude_state is None: status = self.kv_mgr.check_status(self.bootstrap_room) - if status in (KVPoll.Success, KVPoll.Failed): + if status == KVPoll.Failed: + self.begin_failure_quiescence() + if not self.is_transfer_quiesced(): + return KVPoll.Transferring + self.conclude_state = status + self.trace_ctx.trace_req_finish() + elif status == KVPoll.Success: + lifetime = self.kv_mgr._abort_room(self.bootstrap_room) + if lifetime is not None and not lifetime.is_quiesced(): + return KVPoll.Transferring self.conclude_state = status self.trace_ctx.trace_req_finish() elif status == KVPoll.Bootstrapping: timeout_result = self._check_bootstrap_timeout() if timeout_result is not None: - return timeout_result + self.begin_failure_quiescence() + return ( + KVPoll.Failed + if self.is_transfer_quiesced() + else KVPoll.Transferring + ) return status else: return self.conclude_state + def begin_failure_quiescence(self): + if self._failure_quiescence_started: + return + self._failure_quiescence_started = True + self.kv_mgr._abort_room(self.bootstrap_room) + super().abort() + + def is_transfer_quiesced(self) -> bool: + lifetime = self.kv_mgr._find_room_lifetime(self.bootstrap_room) + return lifetime is None or lifetime.is_quiesced() + + def is_failure_quiescing(self) -> bool: + return self._failure_quiescence_started + + def clear(self) -> None: + super().clear() + self.kv_mgr._clear_room_lifetime(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: @@ -1877,7 +2148,7 @@ def _init_trace_ctx(self): self.trace_ctx.trace_req_start() def abort(self): - super().abort() + self.begin_failure_quiescence() self.trace_ctx.abort(abort_info={"reason": "Aborted"}) self.trace_ctx.trace_req_finish() @@ -1889,9 +2160,49 @@ def __init__( bootstrap_addr: str, bootstrap_room: Optional[int] = None, ): - self.session_id = mgr.get_session_id() - self.init_time = None - super().__init__(mgr, bootstrap_addr, bootstrap_room) + with mgr._abort_receivers_lock: + if bootstrap_room is not None and bootstrap_room in mgr.abort_receivers: + raise RuntimeError( + f"Mooncake bootstrap room {bootstrap_room} is already active" + ) + self.session_id = mgr.get_session_id() + self.init_time = None + self._expected_abort_acks = set() + self._received_abort_acks = set() + self._abort_targets = [] + self._abort_ack_lock = threading.Lock() + self._last_abort_send = float("-inf") + self._staging_registered = False + super().__init__(mgr, bootstrap_addr, bootstrap_room) + mgr.abort_receivers[bootstrap_room] = self + if bootstrap_room in mgr.aborting_rooms: + mgr.record_failure( + bootstrap_room, + "Mooncake bootstrap room cannot be reused after abort", + ) + mgr.update_status(bootstrap_room, KVPoll.Failed) + self.conclude_state = KVPoll.Failed + + def init(self, prefill_dp_rank: int): + if self.conclude_state == KVPoll.Failed: + return + super().init(prefill_dp_rank) + + def _get_bootstrap_info_from_server( + self, prefill_dp_rank, prefill_cp_rank, target_tp_rank, target_pp_rank + ): + info = super()._get_bootstrap_info_from_server( + prefill_dp_rank, prefill_cp_rank, target_tp_rank, target_pp_rank + ) + if info is None: + return None + if QUIESCE_CAPABILITY not in info.get("capabilities", []): + self.kv_mgr.record_failure( + self.bootstrap_room, + f"Mooncake prefill rank is missing {QUIESCE_CAPABILITY}", + ) + return None + return info def _register_kv_args(self): for bootstrap_info in self.bootstrap_infos: @@ -1945,6 +2256,7 @@ def _register_kv_args(self): packed_state_dim_per_tensor, packed_staging_base_ptr, staging_total_size_str, + QUIESCE_CAPABILITY_BYTES, ] ) @@ -1968,11 +2280,12 @@ def send_metadata( and self.kv_mgr._staging_ctx.allocator is not None ): self.chunk_staging_infos = [] - self.kv_mgr.register_staging_room_bootstrap( + self._staging_registered = self.kv_mgr.register_staging_room_bootstrap( self.bootstrap_room, self.bootstrap_infos, self ) - for bootstrap_info in self.bootstrap_infos: + abort_targets = self._ensure_abort_targets() + for bootstrap_info, abort_token in abort_targets: sock, lock = self._connect_to_bootstrap_server(bootstrap_info) is_dummy = bootstrap_info["is_dummy"] @@ -1992,24 +2305,119 @@ def send_metadata( ), str(self.required_dst_info_num).encode("ascii"), str(decode_prefix_len or 0).encode("ascii"), + QUIESCE_CAPABILITY_BYTES, + abort_token, ] ) self.init_time = time.time() def poll(self) -> KVPoll: + if self.conclude_state == KVPoll.Failed: + self.begin_failure_quiescence() + if not self.is_transfer_quiesced(): + return KVPoll.Transferring + return self.conclude_state if self.conclude_state is not None: return self.conclude_state status = self.kv_mgr.check_status(self.bootstrap_room) - if status in (KVPoll.Success, KVPoll.Failed): + if status == KVPoll.Failed: + self.begin_failure_quiescence() + if not self.is_transfer_quiesced(): + return KVPoll.Transferring + self.conclude_state = status + elif status == KVPoll.Success: self.conclude_state = status elif status == KVPoll.WaitingForInput: timeout_result = self._check_waiting_timeout() if timeout_result is not None: - return timeout_result + self.begin_failure_quiescence() + return ( + KVPoll.Failed + if self.is_transfer_quiesced() + else KVPoll.Transferring + ) return status + def _record_abort_ack(self, token: Optional[bytes]) -> None: + with self._abort_ack_lock: + if token in self._expected_abort_acks: + self._received_abort_acks.add(token) + + def _abort_acks_complete(self) -> bool: + with self._abort_ack_lock: + return self._expected_abort_acks.issubset(self._received_abort_acks) + + def _ensure_abort_targets(self): + with self._abort_ack_lock: + if not self._abort_targets and self.bootstrap_infos: + self._abort_targets = [ + (bootstrap_info, os.urandom(16).hex().encode("ascii")) + for bootstrap_info in self.bootstrap_infos + ] + self._expected_abort_acks.update( + token for _bootstrap_info, token in self._abort_targets + ) + return list(self._abort_targets) + + def abort(self): + self.kv_mgr.aborting_rooms.add(self.bootstrap_room) + handler = getattr(self.kv_mgr, "_staging_handler", None) + if handler is not None: + handler.begin_abort(self.bootstrap_room) + super().abort() + + def begin_failure_quiescence(self): + self.abort() + self._send_abort_notification() + + def is_failure_quiescing(self) -> bool: + return self.abort_notified + + def is_transfer_quiesced(self) -> bool: + if not self._abort_acks_complete(): + return False + handler = getattr(self.kv_mgr, "_staging_handler", None) + if not self._staging_registered: + return True + if handler is None: + return False + return handler.finalize_room_abort(self.bootstrap_room) + + def _send_abort_notification(self): + now = time.monotonic() + self._ensure_abort_targets() + with self._abort_ack_lock: + if now - self._last_abort_send < ABORT_RETRY_INTERVAL_S: + return + self._last_abort_send = now + targets = [ + (bootstrap_info, token) + for bootstrap_info, token in self._abort_targets + if token not in self._received_abort_acks + ] + + for bootstrap_info, token in targets: + 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"), + token, + ] + ) + except Exception as e: + logger.debug( + "Failed to send abort notification for room %s: %s", + self.bootstrap_room, + e, + ) + def failure_exception(self): if self.conclude_state is None: self.conclude_state = KVPoll.Failed @@ -2025,6 +2433,15 @@ def failure_exception(self): self.bootstrap_room, failure_reason, is_from_another_rank=is_propagated ) + def clear(self) -> None: + super().clear() + handler = getattr(self.kv_mgr, "_staging_handler", None) + if handler is not None: + handler.clear_abort(self.bootstrap_room) + with self.kv_mgr._abort_receivers_lock: + if self.kv_mgr.abort_receivers.get(self.bootstrap_room) is self: + self.kv_mgr.abort_receivers.pop(self.bootstrap_room, None) + class MooncakeKVBootstrapServer(CommonKVBootstrapServer): pass diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 5375b70965b3..61ca2b021ebe 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -526,8 +526,12 @@ def set_kv_buffer_tensors(self, k_buffers: list, v_buffers: list, page_size: int } def register_staging_room_bootstrap(self, room, bootstrap_infos, receiver): + handler = getattr(self, "_staging_handler", None) + if handler is not None: + return handler.register_room_bootstrap(room, bootstrap_infos, receiver) self._staging_ctx.room_bootstrap[room] = bootstrap_infos self._staging_ctx.room_receivers[room] = receiver + return True def _is_watermark_ready( self, agent_name: str, alloc_round: int, alloc_end: int @@ -555,37 +559,21 @@ def decode_staging_thread(): threading.Thread(target=decode_staging_thread, daemon=True).start() def _handle_staging_req(self, msg): - from sglang.srt.disaggregation.common.staging_handler import ( - handle_staging_req, - ) - room = int(msg[1].decode("ascii")) session_id = msg[4].decode("ascii") handler = self._staging_handler assert ( handler is not None ), "STAGING_REQ received before staging handler initialized" - decode_req = handler._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "STAGING_REQ received for unregistered room=%s, skipping", - room, - ) - return - prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size - handle_staging_req( + accepted = handler.handle_staging_req( msg, - self._staging_ctx.allocator, self.kv_args, self.attn_tp_size, - prefill_tp, getattr(self, "kv_buffer_tensors", None), - self._staging_ctx.room_receivers, - self._staging_ctx.room_bootstrap, ) receiver = self._staging_ctx.room_receivers.get(room) - if receiver is not None: + if accepted and receiver is not None: handler.register_wm_subscriber(receiver, session_id) def _prefetch_staging_reqs(self, room: int): diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 5eeb5af68ec9..3bd85f8eae5b 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -782,7 +782,16 @@ def process_disagg_prefill_inflight_queue( else: logger.warning(error_message) req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) - release_kv_cache(req, self.tree_cache) # unlock the tree + if isinstance(req.finished_reason, FINISH_ABORT): + if self.enable_hicache_storage: + self.tree_cache.release_aborted_request(req.rid) + if ( + req.req_pool_idx is not None or self.tree_cache.supports_mamba() + ) and not req.kv_committed_freed: + release_kv_cache(req, self.tree_cache, is_insert=False) + req.pending_bootstrap = False + else: + release_kv_cache(req, self.tree_cache) # unlock the tree if not isinstance(req.finished_reason, FINISH_ABORT): prepare_abort( req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 9ffceda77f94..70f799f47522 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -86,11 +86,20 @@ def _get_failure_prob() -> float: def _poll_with_failure_injection(pollers) -> List[int]: if (failure_prob := _get_failure_prob()) > 0: - return [ + polls = [ int(KVPoll.Failed) if random.random() < failure_prob else int(poller.poll()) for poller in pollers ] - return [int(poller.poll()) for poller in pollers] + else: + polls = [int(poller.poll()) for poller in pollers] + return [ + ( + int(KVPoll.Failed) + if getattr(poller, "is_failure_quiescing", lambda: False)() + else poll + ) + for poller, poll in zip(pollers, polls) + ] def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool: @@ -118,6 +127,34 @@ def _apply_metadata_gate(polls, decode_reqs, metadata_buffers, server_args) -> N polls[i] = int(KVPoll.Transferring) +def _reduce_polls_with_quiescence(pollers, polls, groups): + poll_tensor = torch.tensor(polls, dtype=torch.uint8, device="cpu") + for group in groups: + dist.all_reduce(poll_tensor, op=dist.ReduceOp.MIN, group=group) + + reduced_polls = poll_tensor.tolist() + failed = [poll == int(KVPoll.Failed) for poll in reduced_polls] + if not any(failed): + return reduced_polls + + for poller, is_failed in zip(pollers, failed): + if is_failed: + poller.begin_failure_quiescence() + + quiesced = [ + int(not is_failed or poller.is_transfer_quiesced()) + for poller, is_failed in zip(pollers, failed) + ] + quiesced_tensor = torch.tensor(quiesced, dtype=torch.uint8, device="cpu") + for group in groups: + dist.all_reduce(quiesced_tensor, op=dist.ReduceOp.MIN, group=group) + + for i, is_failed in enumerate(failed): + if is_failed and not quiesced_tensor[i].item(): + reduced_polls[i] = int(KVPoll.Transferring) + return reduced_polls + + def poll_and_all_reduce( pollers, gloo_group: dist.ProcessGroup, @@ -135,9 +172,7 @@ def poll_and_all_reduce( and server_args is not None ): _apply_metadata_gate(polls, decode_reqs, metadata_buffers, server_args) - tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu") - dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group) - return tensor_to_reduce.tolist() + return _reduce_polls_with_quiescence(pollers, polls, [gloo_group]) def poll_and_all_reduce_attn_cp_tp_group( @@ -145,19 +180,12 @@ def poll_and_all_reduce_attn_cp_tp_group( attn_cp_cpu_group: dist.ProcessGroup, attn_tp_cpu_group: dist.ProcessGroup, ): - # First sync across attn-tp ranks so all TP participants for a given (dp, cp) - # shard observe the same status transitions. - polls = poll_and_all_reduce(pollers, attn_tp_cpu_group) - - # Then sync across attn-cp ranks, so all TPxCP participants in one DP shard - # converge to the same global status. - tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu") - dist.all_reduce( - tensor_to_reduce, - op=dist.ReduceOp.MIN, - group=attn_cp_cpu_group, + polls = _poll_with_failure_injection(pollers) + return _reduce_polls_with_quiescence( + pollers, + polls, + [attn_tp_cpu_group, attn_cp_cpu_group], ) - return tensor_to_reduce.tolist() def poll_and_all_reduce_with_staging( @@ -186,9 +214,13 @@ def poll_and_all_reduce_with_staging( # Apply metadata gate on the decode requests to downgrade Success → Transferring for requests whose metadata hasn't landed. if metadata_buffers is not None and server_args is not None: _apply_metadata_gate(raw_polls, decode_reqs, metadata_buffers, server_args) - poll_tensor = torch.tensor(raw_polls, dtype=torch.uint8, device="cpu") - dist.all_reduce(poll_tensor, op=dist.ReduceOp.MIN, group=gloo_group) - return poll_tensor.tolist() + return _reduce_polls_with_quiescence(receivers, raw_polls, [gloo_group]) + + +def defer_chunked_prefill_abort(req, inflight_queue) -> None: + req.disagg_kv_sender.abort() + if req not in inflight_queue: + inflight_queue.append(req) ######################### diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 020cc8ac0fc8..06b00c6ea16b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -61,6 +61,7 @@ MetadataBuffers, ReqToMetadataIdxAllocator, TransferBackend, + defer_chunked_prefill_abort, prepare_abort, ) from sglang.srt.distributed import get_pp_group, get_world_group @@ -2557,21 +2558,18 @@ def process_pending_chunked_abort(self) -> None: req.time_stats.trace_ctx.abort(abort_info={"reason": "Aborted"}) req.to_finish = None if self.disaggregation_mode == DisaggregationMode.PREFILL: - req.disagg_kv_sender.abort() - maybe_release_metadata_buffer( - req, self.req_to_metadata_buffer_idx_allocator - ) - req.pending_bootstrap = False - if self.enable_hicache_storage: - self.tree_cache.release_aborted_request(req.rid) - if ( - req.req_pool_idx is not None or self.tree_cache.supports_mamba() - ) and not req.kv_committed_freed: - release_kv_cache(req, self.tree_cache, is_insert=False) + defer_chunked_prefill_abort(req, self.disagg_prefill_inflight_queue) + else: + if self.enable_hicache_storage: + self.tree_cache.release_aborted_request(req.rid) + if ( + req.req_pool_idx is not None or self.tree_cache.supports_mamba() + ) and not req.kv_committed_freed: + release_kv_cache(req, self.tree_cache, is_insert=False) + self.ipc_channels.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req) self.chunked_req = None self._pending_chunked_abort_req = None - self.ipc_channels.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req) logger.debug(f"Abort chunked prefill request. {req.rid=}") def _build_hisparse_decode_batch(self, reqs):