diff --git a/python/sglang/srt/disaggregation/common/staging_buffer.py b/python/sglang/srt/disaggregation/common/staging_buffer.py index 0824af25a37e..f409c872983b 100644 --- a/python/sglang/srt/disaggregation/common/staging_buffer.py +++ b/python/sglang/srt/disaggregation/common/staging_buffer.py @@ -237,8 +237,17 @@ def free(self, alloc_id: int): self.alloc_order.pop(0) if not self.allocations: + # Once the ring is empty, every byte is safe to reuse. Keeping + # ``head`` at the end of the last allocation leaves the unused + # tail represented as unsafe after the next wrap: an allocation + # in the new round whose end extends past the old head can then + # wait forever for a watermark that can no longer advance. + # Start a fresh round at offset zero so the empty-ring + # watermark covers the whole previous round. + self.round += 1 + self.head = 0 self.watermark_round = self.round - self.watermark_tail = self.head + self.watermark_tail = 0 elif self.alloc_order: off, _, rnd = self.allocations[self.alloc_order[0]] self.watermark_round = rnd diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 525799979e2f..38142cdbeeb5 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -107,8 +107,21 @@ def register_wm_subscriber(self, receiver, session_id: str) -> None: if receiver is None or not receiver.bootstrap_infos: 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) + if key in self._wm_subscribers: + return + + self._wm_subscribers[key] = (receiver, session_id) + # The allocator round is global, while prefills learn watermarks per + # session. A newly registered prefill may therefore receive its first + # allocation after another session has already advanced the ring. Send + # the current watermark immediately; waiting for the next free can + # deadlock when that first allocation is itself waiting on the missed + # watermark. + self._send_watermark( + receiver, + session_id, + self.staging_allocator.get_watermark(), + ) def num_writers_for(self, receiver) -> int: """Compute all TP and PP writers expected for a staging chunk.""" @@ -442,26 +455,28 @@ def _scatter_region( return True def _free_and_send_watermark( - self, alloc_id: int, decode_req: DecodeRequest + 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 - wm_round, wm_tail = post_wm + for receiver, session_id in list(self._wm_subscribers.values()): + self._send_watermark(receiver, session_id, post_wm) + + @staticmethod + def _send_watermark(receiver, session_id: str, watermark) -> None: + """Send one allocator watermark to a registered prefill session.""" + wm_round, wm_tail = watermark 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()): - sid_b = session_id.encode("ascii") - for bootstrap_info in receiver.bootstrap_infos: - try: - sock, lock = receiver._connect_to_bootstrap_server(bootstrap_info) - with lock: - sock.send_multipart( - [b"WATERMARK", wm_round_b, wm_tail_b, sid_b] - ) - except Exception: - pass + sid_b = session_id.encode("ascii") + for bootstrap_info in receiver.bootstrap_infos: + try: + sock, lock = receiver._connect_to_bootstrap_server(bootstrap_info) + with lock: + sock.send_multipart([b"WATERMARK", wm_round_b, wm_tail_b, sid_b]) + except Exception: + pass def is_watermark_ready( diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py index c42f6f65cff9..5bd3b98a2990 100644 --- a/python/sglang/srt/disaggregation/common/utils.py +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -32,6 +32,12 @@ class TransferKVChunk: # Set when the staging worker first counts this chunk toward the per-room # outstanding count; stays set across re-enqueue on a watermark defer. staging_counted: bool = False + # A heterogeneous-TP chunk fans out to multiple decode sessions. Their + # staging watermarks advance independently, so one worker pass can finish + # only a prefix of the fan-out before another destination defers. Keep the + # completed sessions on the re-enqueued chunk: replaying them can write into + # a staging slot that decode has already scattered and recycled. + staging_completed_sessions: set[str] = dataclasses.field(default_factory=set) # Mori early-send: CUDA event to synchronize before RDMA (optional). wait_event: Optional[object] = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index e49da8ee4473..93ccaebc276b 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,12 +21,14 @@ from __future__ import annotations import logging +import queue +import threading import time from collections import deque from concurrent.futures import Future from dataclasses import dataclass from http import HTTPStatus -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple import numpy as np import torch @@ -59,6 +61,8 @@ poll_and_all_reduce_pp, poll_and_all_reduce_with_staging, prepare_abort, + prepare_poll_tensor, + prepare_poll_tensor_with_staging, setup_state_kv_args, ) from sglang.srt.environ import envs @@ -106,6 +110,136 @@ logger = logging.getLogger(__name__) + +class DecodePollCoordinator: + """Run ordered decode collectives outside the scheduler thread. + + Gloo's ``async_op=True`` only makes completion asynchronous; submitting the + collective can still wait for peer ranks. The decode schedulers may be on + different local iterations when attention TP gather is disabled, so even + submission must not run on their main threads. Request broadcasts use the + same worker as status polls: handing one process group between threads after + local completion can let a fast rank enter the next broadcast while a slow + rank is still returning from the preceding all-reduce. + """ + + def __init__(self, group: ProcessGroup): + self._group = group + self._inputs = queue.Queue(maxsize=1) + self._next_sequence = 0 + self._pending_poll = None + self._current_task = None + self._last_stall_log_at = 0.0 + self._thread = threading.Thread( + target=self._run, + name="disagg-decode-poll", + daemon=True, + ) + self._thread.start() + + def submit(self, state: Dict[str, Any]) -> None: + if self._pending_poll is not None: + raise RuntimeError("A disaggregated decode poll is already pending") + task = self._make_task("poll", state) + self._pending_poll = task + self._inputs.put_nowait(task) + + def execute(self, fn: Callable[[], Any]) -> Any: + """Run the next blocking collective on the poll worker.""" + if self._pending_poll is not None: + raise RuntimeError( + "Cannot run a decode collective call before the pending poll " + "result is consumed" + ) + task = self._make_task("call", fn) + self._inputs.put_nowait(task) + task["done"].wait() + return self._unwrap_result(task["result"]) + + def poll(self) -> Optional[Dict[str, Any]]: + task = self._pending_poll + if task is None: + raise RuntimeError("No disaggregated decode poll is pending") + if not task["done"].is_set(): + return None + self._pending_poll = None + return self._unwrap_result(task["result"]) + + def _make_task(self, kind: str, payload: Any) -> Dict[str, Any]: + sequence = self._next_sequence + self._next_sequence += 1 + return { + "sequence": sequence, + "kind": kind, + "payload": payload, + "result": None, + "done": threading.Event(), + "submitted_at": time.monotonic(), + } + + def maybe_log_stall(self, scheduler: Scheduler) -> None: + task = self._pending_poll + if task is None: + return + now = time.monotonic() + age = now - task["submitted_at"] + if age < 10 or now - self._last_stall_log_at < 10: + return + self._last_stall_log_at = now + current = self._current_task + logger.warning( + "Disaggregated decode poll pending for %.1fs: tp_rank=%s " + "poll_seq=%s poll_done=%s worker_task=%s input_qsize=%s " + "scheduler_pending=%s staged_result=%s engine_paused=%s " + "polling_count=%s prealloc=%s transfer=%s retracted=%s", + age, + get_parallel().tp_rank, + task["sequence"], + task["done"].is_set(), + None if current is None else (current["sequence"], current["kind"]), + self._inputs.qsize(), + scheduler.disagg_decode_poll_pending, + scheduler.disagg_decode_poll_result is not None, + scheduler._engine_paused, + getattr(scheduler, "polling_count", None), + len(scheduler.disagg_decode_prealloc_queue.queue), + len(scheduler.disagg_decode_transfer_queue.queue), + len(scheduler.disagg_decode_prealloc_queue.retracted_queue), + ) + + @staticmethod + def _unwrap_result(result: Any) -> Any: + if isinstance(result, BaseException): + raise RuntimeError("Disaggregated decode collective failed") from result + return result + + def _run(self) -> None: + while True: + task = self._inputs.get() + self._current_task = task + kind = task["kind"] + payload = task["payload"] + try: + if kind == "poll": + torch.distributed.all_reduce( + payload["tensor"], + op=torch.distributed.ReduceOp.MIN, + group=self._group, + ) + result = payload + elif kind == "call": + result = payload() + else: + raise RuntimeError(f"Unknown decode collective task {kind!r}") + except BaseException as exc: + task["result"] = exc + task["done"].set() + return + task["result"] = result + task["done"].set() + self._current_task = None + + _is_npu = is_npu() if TYPE_CHECKING: @@ -855,19 +989,19 @@ def _update_handshake_waiters( rids_to_check: Optional[List[str]] = None, pp_good_rids: Optional[List[str]] = None, pp_bad_rids: Optional[List[str]] = None, + precomputed_polls: Optional[Tuple[List[DecodeRequest], List[int]]] = None, ) -> None: - if not self.queue: - return - - # Still poll if any receiver was aborted, otherwise it stays stuck. - if ( - self.pp_size <= 1 - and all(decode_req.waiting_for_input for decode_req in self.queue) - and not any( - decode_req.kv_receiver.conclude_state == KVPoll.Failed - for decode_req in self.queue + if precomputed_polls is not None: + poll_window, polls = precomputed_polls + elif self.pp_size <= 1: + collective_size = self.req_to_metadata_buffer_idx_allocator.size + poll_window = self.queue[:collective_size] + polls = poll_and_all_reduce( + [decode_req.kv_receiver for decode_req in poll_window], + self.gloo_group, + collective_size=collective_size, ) - ): + elif not self.queue: return if self.pp_size > 1: @@ -877,12 +1011,8 @@ def _update_handshake_waiters( pp_good_rids, pp_bad_rids, ) - else: - polls = poll_and_all_reduce( - [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group - ) - - for decode_req, poll in zip(self.queue, polls): + decode_reqs_polled = poll_window if self.pp_size <= 1 else self.queue + for decode_req, poll in zip(decode_reqs_polled, polls): if poll is None: continue if rids_to_check is not None and decode_req.req.rid not in rids_to_check: @@ -916,6 +1046,18 @@ def _update_handshake_waiters( else: raise ValueError(f"Unexpected poll case: {poll}") + def prepare_poll_tensor(self): + """Build this phase's local poll tensor for the combined TP collective.""" + assert self.pp_size <= 1 + self._resolve_pending_reqs() + collective_size = self.req_to_metadata_buffer_idx_allocator.size + poll_window = self.queue[:collective_size] + _, tensor = prepare_poll_tensor( + [decode_req.kv_receiver for decode_req in poll_window], + collective_size=collective_size, + ) + return poll_window, tensor + def _ensure_prefill_info( self, addr_to_reqs: Dict[str, List[DecodeRequest]] ) -> Tuple[Dict[str, List[DecodeRequest]], List[DecodeRequest]]: @@ -1073,6 +1215,7 @@ def pop_preallocated( rids_to_check: Optional[List[str]] = None, pp_good_rids: Optional[List[str]] = None, pp_bad_rids: Optional[List[str]] = None, + precomputed_polls: Optional[Tuple[List[DecodeRequest], List[int]]] = None, ) -> Tuple[List[DecodeRequest], List[DecodeRequest]]: """Pop the preallocated requests from the pending queue (FIFO).""" is_pp_mode = self.pp_size > 1 @@ -1081,8 +1224,14 @@ def pop_preallocated( if is_pp_mode and rids_to_check is not None: raise ValueError("rids_to_check cannot be used in PP mode") - self._resolve_pending_reqs() - self._update_handshake_waiters(rids_to_check, pp_good_rids, pp_bad_rids) + if precomputed_polls is None: + self._resolve_pending_reqs() + self._update_handshake_waiters( + rids_to_check, + pp_good_rids, + pp_bad_rids, + precomputed_polls=precomputed_polls, + ) if is_pp_mode: rids_to_check = set(pp_good_rids) | set(pp_bad_rids) @@ -2218,6 +2367,7 @@ def _poll_with_metadata_gate(self) -> List[int]: self.gloo_group, decode_reqs=self.queue, metadata_buffers=self.metadata_buffers, + collective_size=self.req_to_metadata_buffer_idx_allocator.size, ) def _poll_with_staging(self) -> list: @@ -2226,8 +2376,36 @@ def _poll_with_staging(self) -> list: self.staging_handler, self.gloo_group, metadata_buffers=self.metadata_buffers, + collective_size=self.req_to_metadata_buffer_idx_allocator.size, ) + def prepare_poll_tensor(self): + """Build this phase's local poll tensor for the combined TP collective.""" + if self.scheduler.enable_decode_hicache: + self._process_hicache_local_restores(self.queue) + + collective_size = self.req_to_metadata_buffer_idx_allocator.size + if self.enable_staging: + polls, tensor = prepare_poll_tensor_with_staging( + self.queue, + self.staging_handler, + metadata_buffers=self.metadata_buffers, + collective_size=collective_size, + ) + else: + pollers = ( + [HiCacheRestoreGatedKVReceiver(dr) for dr in self.queue] + if self.scheduler.enable_decode_hicache + else [dr.kv_receiver for dr in self.queue] + ) + polls, tensor = prepare_poll_tensor( + pollers, + decode_reqs=self.queue, + metadata_buffers=self.metadata_buffers, + collective_size=collective_size, + ) + return len(polls), tensor + def _init_staging_handler(self, kv_manager): """Create staging handler from kv_manager. Must be called exactly once.""" from sglang.srt.disaggregation.common.staging_handler import ( @@ -2239,11 +2417,15 @@ def _init_staging_handler(self, kv_manager): ) kv_manager._staging_handler = self.staging_handler - def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req]: - if not self.queue: + def pop_transferred( + self, + rids_to_check: Optional[List[str]] = None, + precomputed_polls: Optional[List[int]] = None, + ) -> List[Req]: + if precomputed_polls is None and not self.queue: return [] - if self.scheduler.enable_decode_hicache: + if precomputed_polls is None and self.scheduler.enable_decode_hicache: self._process_hicache_local_restores( [ decode_req @@ -2252,7 +2434,9 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ] ) - if self.enable_staging: + if precomputed_polls is not None: + polls = precomputed_polls + elif self.enable_staging: polls = self._poll_with_staging() else: polls = self._poll_with_metadata_gate() @@ -2448,12 +2632,25 @@ def event_loop_normal_disagg_decode(self: Scheduler): """A normal scheduler loop for decode worker in disaggregation mode.""" while True: + # Do not enter the next request broadcast until this rank's + # background decode-poll epoch has completed. Otherwise a fast + # scheduler rank can block in request ingress while a peer's poll + # thread still needs that rank to join the collective. + if not self._stage_completed_decode_poll(): + continue + # Pending rooms from the prior cycle can overlap request intake and # the tail of the in-flight decode graph. if not self._engine_paused: self.disagg_decode_prealloc_queue.prefetch_prefill_dp_rank_queries() # Receive requests - recv_reqs = self.request_receiver.recv_requests() + recv_reqs = self.request_receiver.recv_requests( + collective_executor=( + self.disagg_decode_poll_coordinator.execute + if self.disagg_decode_prealloc_queue.pp_size == 1 + else None + ) + ) self.process_input_requests(recv_reqs) if self._engine_paused: continue @@ -2491,12 +2688,24 @@ def pop_and_process(): self.process_batch_result(tmp_batch, tmp_result) while True: + # Keep request ingress and decode-poll epochs in the same order on + # every rank even when attention TP gather is disabled and the + # scheduler loops advance at different rates. + if not self._stage_completed_decode_poll(): + continue + # Pending rooms from the prior cycle can overlap request intake and # the tail of the in-flight decode graph. if not self._engine_paused: self.disagg_decode_prealloc_queue.prefetch_prefill_dp_rank_queries() # Receive requests - recv_reqs = self.request_receiver.recv_requests() + recv_reqs = self.request_receiver.recv_requests( + collective_executor=( + self.disagg_decode_poll_coordinator.execute + if self.disagg_decode_prealloc_queue.pp_size == 1 + else None + ) + ) self.process_input_requests(recv_reqs) if self._engine_paused: continue @@ -2659,6 +2868,56 @@ def get_new_prebuilt_batch( return new_batch + def _stage_completed_decode_poll(self: Scheduler) -> bool: + """Stage the pending poll result before the next request broadcast. + + Collective submission runs on ``DecodePollCoordinator`` so the + scheduler thread never blocks in Gloo. This gate keeps the scheduler + out of request ingress until that already-submitted epoch is complete, + preventing rank-skewed loops from splitting across the two operations. + """ + if self.disagg_decode_prealloc_queue.pp_size > 1: + return True + if getattr(self, "disagg_decode_poll_result", None) is not None: + return True + if not self.disagg_decode_poll_pending: + return True + + poll_state = self.disagg_decode_poll_coordinator.poll() + if poll_state is None: + self.disagg_decode_poll_coordinator.maybe_log_stall(self) + return False + + self.disagg_decode_poll_result = poll_state + self.disagg_decode_poll_pending = False + return True + + def _submit_decode_poll(self: Scheduler, has_retracted_reqs: bool) -> None: + prealloc_window, prealloc_tensor = ( + self.disagg_decode_prealloc_queue.prepare_poll_tensor() + ) + transfer_count, transfer_tensor = ( + self.disagg_decode_transfer_queue.prepare_poll_tensor() + ) + can_process_tensor = torch.tensor( + [not has_retracted_reqs], dtype=torch.uint8, device="cpu" + ) + combined_tensor = torch.cat( + ( + prealloc_tensor, + transfer_tensor, + can_process_tensor, + ) + ) + self.disagg_decode_poll_coordinator.submit( + { + "tensor": combined_tensor, + "prealloc_window": prealloc_window, + "transfer_count": transfer_count, + } + ) + self.disagg_decode_poll_pending = True + def process_decode_queue(self: Scheduler): if self.enable_decode_hicache: self.tree_cache.check_hicache_events() @@ -2673,7 +2932,8 @@ def process_decode_queue(self: Scheduler): # try to resume retracted requests if there are enough space for another `num_reserved_decode_tokens` decode steps resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() self.waiting_queue.extend(resumed_reqs) - if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: + has_retracted_reqs = len(self.disagg_decode_prealloc_queue.retracted_queue) > 0 + if self.disagg_decode_prealloc_queue.pp_size > 1 and has_retracted_reqs: # if there are still retracted requests, we do not allocate new requests return @@ -2684,11 +2944,67 @@ def process_decode_queue(self: Scheduler): self.polling_count = (self.polling_count + 1) % self.polling_interval if self.polling_count % self.polling_interval == 0: - req_conns, _ = self.disagg_decode_prealloc_queue.pop_preallocated() - self.disagg_decode_transfer_queue.extend(req_conns) - transferred_reqs = ( - self.disagg_decode_transfer_queue.pop_transferred() - ) # the requests which kv has arrived + if self.disagg_decode_prealloc_queue.pp_size > 1: + req_conns, _ = self.disagg_decode_prealloc_queue.pop_preallocated() + self.disagg_decode_transfer_queue.extend(req_conns) + transferred_reqs = self.disagg_decode_transfer_queue.pop_transferred() + else: + completed_poll = False + poll_state = getattr(self, "disagg_decode_poll_result", None) + if poll_state is not None: + self.disagg_decode_poll_result = None + completed_poll = True + elif self.disagg_decode_poll_pending: + poll_state = self.disagg_decode_poll_coordinator.poll() + if poll_state is None: + return + self.disagg_decode_poll_pending = False + completed_poll = True + + if completed_poll: + combined_tensor = poll_state["tensor"] + prealloc_window = poll_state["prealloc_window"] + transfer_count = poll_state["transfer_count"] + + collective_size = self.req_to_metadata_buffer_idx_allocator.size + can_process_reqs = bool(combined_tensor[-1].item()) + if can_process_reqs: + prealloc_polls = combined_tensor[ + : len(prealloc_window) + ].tolist() + transfer_polls = combined_tensor[ + collective_size : collective_size + transfer_count + ].tolist() + + # Consume the transfer snapshot before adding requests + # that completed preallocation in this epoch. New + # transfers are included in the next poll. + transferred_reqs = ( + self.disagg_decode_transfer_queue.pop_transferred( + precomputed_polls=transfer_polls + ) + ) + req_conns, _ = ( + self.disagg_decode_prealloc_queue.pop_preallocated( + precomputed_polls=( + prealloc_window, + prealloc_polls, + ) + ) + ) + self.disagg_decode_transfer_queue.extend(req_conns) + else: + # Preserve the existing retraction gate on every rank. + transferred_reqs = [] + + # disable-attn-tp-gather permits scheduler ranks to consume a + # completed epoch on different local ticks. Refill the + # coordinator in the same invocation that drains its result; + # otherwise a rank can enter the next request broadcast while + # its peers already wait for that rank in the next poll. + self._submit_decode_poll(has_retracted_reqs) + if not completed_poll: + return 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 1c0955653e8c..aba2667cef26 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1628,13 +1628,11 @@ def transfer_worker( and staging_buffer is not None ): staging_strategy = self._try_create_staging_strategy(staging_buffer) - reqs_to_be_processed = ( + reqs_to_be_processed = list( self.transfer_infos[kv_chunk.room].values() if kv_chunk.room in self.transfer_infos else [] ) - 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) @@ -1647,6 +1645,11 @@ def transfer_worker( for req in reqs_to_be_processed: start_ts = time.perf_counter() if not req.is_dummy: + if ( + req.mooncake_session_id + in kv_chunk.staging_completed_sessions + ): + continue # Early exit if the request has failed with self.session_lock: if req.mooncake_session_id in self.failed_sessions: @@ -1848,23 +1851,26 @@ def transfer_worker( kv_chunk.prefill_aux_index, target_rank_registration_info.dst_aux_ptrs, ) - polls.append(True if ret == 0 else False) - dst_ranks_infos.append( - (req.endpoint, req.dst_port, req.room) - ) + if ret != 0: + with self.session_lock: + self.session_failures[req.mooncake_session_id] += 1 + self.failed_sessions.add(req.mooncake_session_id) + self.record_failure( + kv_chunk.room, + f"Failed to send aux data of {kv_chunk.room} to " + f"{NetworkAddress(req.endpoint, req.dst_port).to_host_port_str()}", + ) + self.update_status(kv_chunk.room, KVPoll.Failed) + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + KVPoll.Failed, + prefill_unique_rank, + ) + break - # Only sync status when all the dst ranks have received the kvcache - if len(polls) == req.required_dst_info_num: - status = KVPoll.Success if all(polls) else KVPoll.Failed - self.update_status(req.room, status) - for endpoint, dst_port, room in dst_ranks_infos: - self.sync_status_to_decode_endpoint( - endpoint, - dst_port, - room, - status, - prefill_unique_rank, - ) + kv_chunk.staging_completed_sessions.add(req.mooncake_session_id) else: # Dummy request means the decode instance is not used, so its status can be marked as success directly # Dummy request does not need to sync status to decode endpoint @@ -1888,6 +1894,24 @@ def transfer_worker( if staging_deferred: continue + if ( + kv_chunk.is_last_chunk + and self.check_status(kv_chunk.room) != KVPoll.Failed + ): + non_dummy_reqs = [ + req for req in reqs_to_be_processed if not req.is_dummy + ] + if len(kv_chunk.staging_completed_sessions) == len(non_dummy_reqs): + self.update_status(kv_chunk.room, KVPoll.Success) + for req in non_dummy_reqs: + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + KVPoll.Success, + prefill_unique_rank, + ) + self._staging_outstanding[kv_chunk.room] -= 1 if self.enable_deferred_decode_kv_release: # In-flight write finished; if aborted and nothing outstanding, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 09fd8efa944e..bc6b1e0c663b 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -207,14 +207,52 @@ def poll_and_all_reduce( gloo_group: dist.ProcessGroup, decode_reqs=None, metadata_buffers: Optional[MetadataBuffers] = None, + collective_size: Optional[int] = None, ): - # at a certain prob, the poll is failed to simulate failure - polls = _poll_with_failure_injection(pollers) + polls, tensor_to_reduce = prepare_poll_tensor( + pollers, + decode_reqs=decode_reqs, + metadata_buffers=metadata_buffers, + collective_size=collective_size, + ) + dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group) + return tensor_to_reduce[: len(polls)].tolist() - # Apply metadata gate on the decode requests to downgrade Success → Transferring for requests whose metadata hasn't landed. + +def prepare_poll_tensor( + pollers, + decode_reqs=None, + metadata_buffers: Optional[MetadataBuffers] = None, + collective_size: Optional[int] = None, +): + """Poll locally and build a fixed-shape tensor without synchronizing.""" + polls = _poll_with_failure_injection(pollers) if decode_reqs is not None and metadata_buffers is not None: _apply_metadata_gate(polls, decode_reqs, metadata_buffers) - return _all_reduce_polls(polls, gloo_group) + return polls, _build_poll_tensor(polls, collective_size) + + +def _build_poll_tensor(polls: List[int], collective_size: Optional[int]): + if collective_size is None: + return torch.tensor(polls, dtype=torch.uint8, device="cpu") + if len(polls) > collective_size: + raise RuntimeError( + f"PD decode poll queue exceeded its collective capacity: " + f"{len(polls)} > {collective_size}" + ) + + # Every TP rank must issue the same-shaped collective even while its local + # queue is still catching up. Missing entries stay at Bootstrapping so they + # cannot make a peer advance the corresponding request prematurely. + tensor = torch.full( + (collective_size,), + int(KVPoll.Bootstrapping), + dtype=torch.uint8, + device="cpu", + ) + if polls: + tensor[: len(polls)] = torch.tensor(polls, dtype=torch.uint8, device="cpu") + return tensor def poll_and_all_reduce_attn_cp_tp_group( @@ -236,8 +274,26 @@ def poll_and_all_reduce_with_staging( staging_handler, gloo_group: dist.ProcessGroup, metadata_buffers: Optional[MetadataBuffers] = None, + collective_size: Optional[int] = None, ): """Staging-aware polling: advance scatter, demote incomplete transfers, all_reduce.""" + raw_polls, poll_tensor = prepare_poll_tensor_with_staging( + decode_reqs, + staging_handler, + metadata_buffers=metadata_buffers, + collective_size=collective_size, + ) + dist.all_reduce(poll_tensor, op=dist.ReduceOp.MIN, group=gloo_group) + return poll_tensor[: len(raw_polls)].tolist() + + +def prepare_poll_tensor_with_staging( + decode_reqs, + staging_handler, + metadata_buffers: Optional[MetadataBuffers] = None, + collective_size: Optional[int] = None, +): + """Run staging progress and build its local fixed-shape poll tensor.""" for decode_req in decode_reqs: if decode_req.kv_receiver.require_staging and not staging_handler.is_done( decode_req @@ -263,7 +319,8 @@ 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: _apply_metadata_gate(raw_polls, decode_reqs, metadata_buffers) - return _all_reduce_polls(raw_polls, gloo_group) + poll_tensor = _build_poll_tensor(raw_polls, collective_size) + return raw_polls, poll_tensor ######################### diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4250cbe9311a..3ef112dc6093 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -72,6 +72,7 @@ from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix from sglang.srt.disaggregation.decode import ( + DecodePollCoordinator, DecodePreallocQueue, DecodeTransferQueue, SchedulerDisaggregationDecodeMixin, @@ -1393,6 +1394,22 @@ def init_disaggregation(self): if ( self.disaggregation_mode == DisaggregationMode.DECODE ): # *8 headroom for MiniMax-M3; *2 for other models. + self.disagg_decode_poll_pending = False + self.disagg_decode_poll_result = None + # Reuse the request-ingress Gloo group for decode polling. The + # event-loop gate waits for each poll before the next request + # broadcast, so all ranks issue both operations in one order. A + # separate group is unsafe here: a fast rank can submit its poll + # while a slow rank is still returning from the prior broadcast, + # creating a cross-group Gloo progress deadlock. + decode_coordination_group = ( + self.attn_tp_cpu_group + if self.enable_dp_attention + else self.tp_cpu_group + ) + self.disagg_decode_poll_coordinator = DecodePollCoordinator( + decode_coordination_group + ) buffer_multiplier = ( 8 if is_minimax_sparse(self.model_config.hf_config) else 2 ) @@ -1410,7 +1427,7 @@ def init_disaggregation(self): # The decode requests polling kv cache self.disagg_decode_transfer_queue = DecodeTransferQueue( - gloo_group=self.attn_tp_cpu_group, + gloo_group=decode_coordination_group, req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator, tp_rank=self.ps.tp_rank, metadata_buffers=self.disagg_metadata_buffers, @@ -1428,7 +1445,7 @@ def init_disaggregation(self): scheduler=self, transfer_queue=self.disagg_decode_transfer_queue, tree_cache=self.tree_cache, - gloo_group=self.attn_tp_cpu_group, + gloo_group=decode_coordination_group, tp_rank=self.ps.tp_rank, tp_size=self.ps.tp_size, dp_size=get_parallel().dp_size, diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index bcdfff5a336d..86ae50720647 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -75,6 +75,8 @@ def recv_limit_reached(self, num_recv_reqs: int) -> bool: @scheduler_nvtx_method("scheduler.recv_requests") def recv_requests( self, + *, + collective_executor: Optional[Callable[[Callable[[], Any]], Any]] = None, ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: """Receive results at tp_rank = 0 and broadcast it to all other TP ranks.""" @@ -90,7 +92,12 @@ def recv_requests( if self.input_blocker is not None: recv_reqs = self.input_blocker.handle(recv_reqs) - recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs) + if collective_executor is None: + recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs) + else: + recv_reqs = collective_executor( + lambda: self._broadcast_reqs_across_ranks(recv_reqs) + ) if self.ps.pp_rank == 0: self.unwrap_pickle_wrapper(recv_reqs) diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 7b5b9f8a06a0..8f581e37ffa8 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -1,15 +1,24 @@ +import threading +import time import unittest from concurrent.futures import Future from types import SimpleNamespace from unittest.mock import MagicMock, patch +import torch + from sglang.srt.disaggregation.base import KVPoll from sglang.srt.disaggregation.decode import ( + DecodePollCoordinator, DecodePreallocQueue, DecodeTransferQueue, HiCacheRestoreResult, + SchedulerDisaggregationDecodeMixin, +) +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + _build_poll_tensor, ) -from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.managers.schedule_batch import FINISH_ABORT from sglang.srt.managers.scheduler import Scheduler @@ -31,7 +40,428 @@ def failure_exception(self): return None +class FakeDecodeScheduler(SchedulerDisaggregationDecodeMixin, SimpleNamespace): + """Minimal scheduler whose mixin helpers remain real in queue tests.""" + + class TestDecodeQueueCleanup(CustomTestCase): + def test_fixed_poll_tensor_pads_missing_tp_entries_as_bootstrapping(self): + tensor = _build_poll_tensor([KVPoll.Success], collective_size=3) + + self.assertEqual( + tensor.tolist(), + [KVPoll.Success, KVPoll.Bootstrapping, KVPoll.Bootstrapping], + ) + + def test_fixed_poll_tensor_rejects_queue_overflow(self): + with self.assertRaisesRegex(RuntimeError, "exceeded its collective capacity"): + _build_poll_tensor( + [KVPoll.Bootstrapping, KVPoll.WaitingForInput], collective_size=1 + ) + + @patch("sglang.srt.disaggregation.decode.poll_and_all_reduce") + def test_prealloc_poll_uses_metadata_sized_fifo_window(self, mock_poll): + receivers = [MagicMock() for _ in range(3)] + decode_reqs = [ + SimpleNamespace( + kv_receiver=receiver, + waiting_for_input=False, + req=SimpleNamespace(time_stats=MagicMock()), + ) + for receiver in receivers + ] + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.pp_size = 1 + queue.queue = decode_reqs + queue.gloo_group = MagicMock() + queue.req_to_metadata_buffer_idx_allocator = SimpleNamespace(size=2) + mock_poll.return_value = [KVPoll.WaitingForInput] * 2 + + queue._update_handshake_waiters() + + mock_poll.assert_called_once_with( + receivers[:2], queue.gloo_group, collective_size=2 + ) + self.assertTrue(decode_reqs[0].waiting_for_input) + self.assertTrue(decode_reqs[1].waiting_for_input) + self.assertFalse(decode_reqs[2].waiting_for_input) + + def test_empty_transfer_queue_builds_fixed_combined_poll_segment(self): + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.queue = [] + queue.enable_staging = False + queue.req_to_metadata_buffer_idx_allocator = SimpleNamespace(size=8) + queue.metadata_buffers = MagicMock() + queue.scheduler = SimpleNamespace( + enable_decode_hicache=False, + server_args=MagicMock(), + ) + + count, tensor = queue.prepare_poll_tensor() + + self.assertEqual(count, 0) + self.assertEqual(tensor.tolist(), [KVPoll.Bootstrapping] * 8) + + @patch("sglang.srt.disaggregation.decode.torch.distributed.all_reduce") + def test_poll_coordinator_keeps_collective_off_caller_thread(self, mock_all_reduce): + entered = threading.Event() + release = threading.Event() + + def blocking_all_reduce(*_args, **_kwargs): + entered.set() + self.assertTrue(release.wait(timeout=1)) + + mock_all_reduce.side_effect = blocking_all_reduce + coordinator = DecodePollCoordinator(MagicMock()) + state = {"tensor": torch.tensor([KVPoll.Success], dtype=torch.uint8)} + + coordinator.submit(state) + + self.assertTrue(entered.wait(timeout=1)) + self.assertIsNone(coordinator.poll()) + release.set() + deadline = time.monotonic() + 1 + result = None + while result is None and time.monotonic() < deadline: + result = coordinator.poll() + time.sleep(0.001) + self.assertIs(result, state) + + @patch("sglang.srt.disaggregation.decode.torch.distributed.all_reduce") + def test_poll_and_request_collectives_use_one_worker_thread(self, mock_all_reduce): + collective_threads = [] + + def record_poll(*_args, **_kwargs): + collective_threads.append(threading.get_ident()) + + mock_all_reduce.side_effect = record_poll + coordinator = DecodePollCoordinator(MagicMock()) + state = {"tensor": torch.tensor([KVPoll.Success], dtype=torch.uint8)} + + coordinator.submit(state) + deadline = time.monotonic() + 1 + while coordinator.poll() is None and time.monotonic() < deadline: + time.sleep(0.001) + + result = coordinator.execute( + lambda: collective_threads.append(threading.get_ident()) or ["request"] + ) + + self.assertEqual(result, ["request"]) + self.assertEqual(len(collective_threads), 2) + self.assertEqual(collective_threads[0], collective_threads[1]) + self.assertNotEqual(collective_threads[0], threading.get_ident()) + + @patch("sglang.srt.disaggregation.decode.torch.distributed.all_reduce") + def test_poll_result_is_bound_to_its_task(self, mock_all_reduce): + entered = threading.Event() + release = threading.Event() + + def blocking_all_reduce(*_args, **_kwargs): + entered.set() + self.assertTrue(release.wait(timeout=1)) + + mock_all_reduce.side_effect = blocking_all_reduce + coordinator = DecodePollCoordinator(MagicMock()) + state = {"tensor": torch.tensor([KVPoll.Success], dtype=torch.uint8)} + + coordinator.submit(state) + self.assertTrue(entered.wait(timeout=1)) + with self.assertRaisesRegex(RuntimeError, "pending poll result"): + coordinator.execute(lambda: []) + with self.assertRaisesRegex(RuntimeError, "already pending"): + coordinator.submit(state) + + release.set() + deadline = time.monotonic() + 1 + result = None + while result is None and time.monotonic() < deadline: + result = coordinator.poll() + time.sleep(0.001) + self.assertIs(result, state) + self.assertEqual(coordinator.execute(lambda: ["request"]), ["request"]) + + @patch("sglang.srt.disaggregation.decode.get_disagg") + def test_decode_queue_submits_one_combined_poll(self, mock_get_disagg): + mock_get_disagg.return_value = SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + disaggregation_decode_polling_interval=1, + ) + prealloc_req = MagicMock() + prealloc_queue = MagicMock(pp_size=1, queue=[prealloc_req], retracted_queue=[]) + prealloc_queue.resume_retracted_reqs.return_value = [] + prealloc_queue.prepare_poll_tensor.return_value = ( + [prealloc_req], + torch.tensor([KVPoll.WaitingForInput, KVPoll.Bootstrapping]), + ) + transfer_queue = MagicMock(queue=[]) + transfer_queue.prepare_poll_tensor.return_value = ( + 1, + torch.tensor([KVPoll.Success, KVPoll.Bootstrapping]), + ) + coordinator = MagicMock() + scheduler = FakeDecodeScheduler( + enable_decode_hicache=False, + disagg_decode_prealloc_queue=prealloc_queue, + disagg_decode_transfer_queue=transfer_queue, + disagg_decode_poll_pending=False, + disagg_decode_poll_coordinator=coordinator, + waiting_queue=[], + ) + + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + + poll_state = coordinator.submit.call_args.args[0] + self.assertEqual( + poll_state["tensor"].tolist(), + [ + KVPoll.WaitingForInput, + KVPoll.Bootstrapping, + KVPoll.Success, + KVPoll.Bootstrapping, + 1, + ], + ) + self.assertIs(poll_state["prealloc_window"][0], prealloc_req) + self.assertEqual(poll_state["transfer_count"], 1) + self.assertTrue(scheduler.disagg_decode_poll_pending) + transfer_queue.pop_transferred.assert_not_called() + prealloc_queue.pop_preallocated.assert_not_called() + + @patch("sglang.srt.disaggregation.decode.get_disagg") + def test_decode_queue_does_not_block_on_incomplete_background_poll( + self, mock_get_disagg + ): + mock_get_disagg.return_value = SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + disaggregation_decode_polling_interval=1, + ) + prealloc_queue = MagicMock(pp_size=1, queue=[], retracted_queue=[]) + prealloc_queue.resume_retracted_reqs.return_value = [] + transfer_queue = MagicMock(queue=[]) + coordinator = MagicMock() + coordinator.poll.return_value = None + scheduler = FakeDecodeScheduler( + enable_decode_hicache=False, + disagg_decode_prealloc_queue=prealloc_queue, + disagg_decode_transfer_queue=transfer_queue, + disagg_decode_poll_pending=True, + disagg_decode_poll_coordinator=coordinator, + waiting_queue=[], + ) + + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + + coordinator.poll.assert_called_once_with() + coordinator.submit.assert_not_called() + prealloc_queue.prepare_poll_tensor.assert_not_called() + transfer_queue.prepare_poll_tensor.assert_not_called() + + def test_pending_poll_blocks_next_request_epoch_until_result_is_staged(self): + coordinator = MagicMock() + coordinator.poll.side_effect = [None, {"tensor": torch.tensor([1])}] + scheduler = FakeDecodeScheduler( + disagg_decode_prealloc_queue=SimpleNamespace(pp_size=1), + disagg_decode_poll_pending=True, + disagg_decode_poll_result=None, + disagg_decode_poll_coordinator=coordinator, + ) + + self.assertFalse( + SchedulerDisaggregationDecodeMixin._stage_completed_decode_poll(scheduler) + ) + self.assertTrue(scheduler.disagg_decode_poll_pending) + self.assertIsNone(scheduler.disagg_decode_poll_result) + + self.assertTrue( + SchedulerDisaggregationDecodeMixin._stage_completed_decode_poll(scheduler) + ) + self.assertFalse(scheduler.disagg_decode_poll_pending) + self.assertEqual(scheduler.disagg_decode_poll_result["tensor"].tolist(), [1]) + self.assertEqual(coordinator.poll.call_count, 2) + + @patch("sglang.srt.disaggregation.decode.get_disagg") + def test_decode_queue_consumes_completed_background_poll(self, mock_get_disagg): + mock_get_disagg.return_value = SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + disaggregation_decode_polling_interval=1, + ) + prealloc_req = MagicMock() + prealloc_queue = MagicMock(pp_size=1, queue=[prealloc_req], retracted_queue=[]) + prealloc_queue.resume_retracted_reqs.return_value = [] + prealloc_queue.pop_preallocated.return_value = (["new-transfer"], []) + prealloc_queue.prepare_poll_tensor.return_value = ( + [], + torch.tensor([KVPoll.Bootstrapping, KVPoll.Bootstrapping]), + ) + transfer_queue = MagicMock(queue=[]) + transfer_queue.pop_transferred.return_value = ["ready"] + transfer_queue.prepare_poll_tensor.return_value = ( + 0, + torch.tensor([KVPoll.Bootstrapping, KVPoll.Bootstrapping]), + ) + coordinator = MagicMock() + coordinator.poll.return_value = { + "tensor": torch.tensor( + [ + KVPoll.WaitingForInput, + KVPoll.Bootstrapping, + KVPoll.Success, + KVPoll.Bootstrapping, + 1, + ], + dtype=torch.uint8, + ), + "prealloc_window": [prealloc_req], + "transfer_count": 1, + } + scheduler = FakeDecodeScheduler( + enable_decode_hicache=False, + enable_hisparse=False, + disagg_decode_prealloc_queue=prealloc_queue, + disagg_decode_transfer_queue=transfer_queue, + req_to_metadata_buffer_idx_allocator=SimpleNamespace(size=2), + disagg_decode_poll_pending=True, + disagg_decode_poll_coordinator=coordinator, + waiting_queue=[], + ) + + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + + self.assertTrue(scheduler.disagg_decode_poll_pending) + coordinator.submit.assert_called_once() + next_poll_state = coordinator.submit.call_args.args[0] + self.assertEqual( + next_poll_state["tensor"].tolist(), + [ + KVPoll.Bootstrapping, + KVPoll.Bootstrapping, + KVPoll.Bootstrapping, + KVPoll.Bootstrapping, + 1, + ], + ) + self.assertEqual(next_poll_state["prealloc_window"], []) + self.assertEqual(next_poll_state["transfer_count"], 0) + transfer_queue.pop_transferred.assert_called_once_with( + precomputed_polls=[KVPoll.Success] + ) + prealloc_queue.pop_preallocated.assert_called_once_with( + precomputed_polls=([prealloc_req], [KVPoll.WaitingForInput]) + ) + transfer_queue.extend.assert_called_once_with(["new-transfer"]) + self.assertEqual(scheduler.waiting_queue, ["ready"]) + + @patch("sglang.srt.disaggregation.decode.get_disagg") + def test_retracted_queue_keeps_fixed_poll_order(self, mock_get_disagg): + mock_get_disagg.return_value = SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + disaggregation_decode_polling_interval=1, + ) + prealloc_queue = MagicMock(pp_size=1, queue=[], retracted_queue=[MagicMock()]) + prealloc_queue.resume_retracted_reqs.return_value = [] + prealloc_queue.prepare_poll_tensor.return_value = ( + [], + torch.tensor([KVPoll.Bootstrapping]), + ) + transfer_queue = MagicMock(queue=[]) + transfer_queue.prepare_poll_tensor.return_value = ( + 0, + torch.tensor([KVPoll.Bootstrapping]), + ) + coordinator = MagicMock() + scheduler = FakeDecodeScheduler( + enable_decode_hicache=False, + disagg_decode_prealloc_queue=prealloc_queue, + disagg_decode_transfer_queue=transfer_queue, + disagg_decode_poll_pending=False, + disagg_decode_poll_coordinator=coordinator, + waiting_queue=[], + ) + + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + + self.assertEqual( + coordinator.submit.call_args.args[0]["tensor"].tolist(), + [KVPoll.Bootstrapping, KVPoll.Bootstrapping, 0], + ) + prealloc_queue.pop_preallocated.assert_not_called() + transfer_queue.pop_transferred.assert_not_called() + + @patch("sglang.srt.disaggregation.decode.get_disagg") + def test_idle_decode_queue_still_submits_poll_epoch(self, mock_get_disagg): + mock_get_disagg.return_value = SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + disaggregation_decode_polling_interval=1, + ) + prealloc_queue = MagicMock(pp_size=1, queue=[], retracted_queue=[]) + prealloc_queue.resume_retracted_reqs.return_value = [] + prealloc_queue.prepare_poll_tensor.return_value = ( + [], + torch.tensor([KVPoll.Bootstrapping]), + ) + transfer_queue = MagicMock(queue=[]) + transfer_queue.prepare_poll_tensor.return_value = ( + 0, + torch.tensor([KVPoll.Bootstrapping]), + ) + coordinator = MagicMock() + scheduler = FakeDecodeScheduler( + enable_decode_hicache=False, + disagg_decode_prealloc_queue=prealloc_queue, + disagg_decode_transfer_queue=transfer_queue, + disagg_decode_poll_pending=False, + disagg_decode_poll_coordinator=coordinator, + waiting_queue=[], + ) + + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + + self.assertEqual( + coordinator.submit.call_args.args[0]["tensor"].tolist(), + [KVPoll.Bootstrapping, KVPoll.Bootstrapping, 1], + ) + coordinator.poll.assert_not_called() + self.assertTrue(scheduler.disagg_decode_poll_pending) + + @patch("sglang.srt.disaggregation.decode.get_disagg") + def test_last_empty_poll_immediately_submits_next_epoch(self, mock_get_disagg): + mock_get_disagg.return_value = SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + disaggregation_decode_polling_interval=1, + ) + prealloc_queue = MagicMock(pp_size=1, queue=[], retracted_queue=[]) + prealloc_queue.resume_retracted_reqs.return_value = [] + prealloc_queue.prepare_poll_tensor.return_value = ( + [], + torch.tensor([KVPoll.Bootstrapping]), + ) + prealloc_queue.pop_preallocated.return_value = ([], []) + transfer_queue = MagicMock(queue=[]) + transfer_queue.prepare_poll_tensor.return_value = ( + 0, + torch.tensor([KVPoll.Bootstrapping]), + ) + transfer_queue.pop_transferred.return_value = [] + coordinator = MagicMock() + scheduler = FakeDecodeScheduler( + enable_decode_hicache=False, + enable_hisparse=False, + disagg_decode_prealloc_queue=prealloc_queue, + disagg_decode_transfer_queue=transfer_queue, + req_to_metadata_buffer_idx_allocator=SimpleNamespace(size=1), + disagg_decode_poll_pending=False, + disagg_decode_poll_coordinator=coordinator, + waiting_queue=[], + ) + + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + coordinator.poll.return_value = coordinator.submit.call_args.args[0] + SchedulerDisaggregationDecodeMixin.process_decode_queue(scheduler) + + self.assertTrue(scheduler.disagg_decode_poll_pending) + self.assertEqual(coordinator.submit.call_count, 2) + def test_paged_swa_retraction_resume_uses_physical_page_budget(self): # resume_retracted_reqs reads the retraction backend off the disagg # bag, so the case publishes a config instead of injecting one. diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 654d1d864123..e3ec62cfe57d 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -1,18 +1,25 @@ import struct import threading import unittest +from collections import defaultdict from types import SimpleNamespace from unittest.mock import Mock, patch import numpy as np import torch -from sglang.srt.disaggregation.base.conn import KVArgs, StateType +from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType from sglang.srt.disaggregation.common.conn import CommonKVManager +from sglang.srt.disaggregation.common.staging_buffer import ( + StagingAllocator, +) from sglang.srt.disaggregation.common.staging_handler import ( + DecodeStagingHandler, handle_staging_req, + is_watermark_ready, ) from sglang.srt.disaggregation.common.utils import ( + TransferKVChunk, group_concurrent_contiguous, pack_int_lists, pack_list_of_buffers, @@ -221,6 +228,172 @@ def test_mismatched_nonempty_lengths_raise(self): class TestMooncakePPStaging(unittest.TestCase): + def test_new_watermark_subscriber_receives_current_allocator_state(self): + sock = Mock() + bootstrap_info = {"host": "prefill", "port": 7200} + receiver = SimpleNamespace( + bootstrap_infos=[bootstrap_info], + _connect_to_bootstrap_server=Mock(return_value=(sock, threading.Lock())), + ) + handler = object.__new__(DecodeStagingHandler) + handler.staging_allocator = SimpleNamespace( + get_watermark=Mock(return_value=(3, 0)) + ) + handler._wm_subscribers = {} + + handler.register_wm_subscriber(receiver, "session-new") + + receiver._connect_to_bootstrap_server.assert_called_once_with(bootstrap_info) + sock.send_multipart.assert_called_once_with( + [b"WATERMARK", b"3", b"0", b"session-new"] + ) + + # Re-registering the same bootstrap does not duplicate subscriptions + # or send an older snapshot over a live stream. + handler.register_wm_subscriber(receiver, "session-new") + self.assertEqual(sock.send_multipart.call_count, 1) + + def test_empty_ring_reset_makes_wrap_gap_reusable(self): + allocator = object.__new__(StagingAllocator) + allocator.total_size = 100 + allocator.head = 0 + allocator.round = 0 + allocator.allocations = {} + allocator.alloc_order = [] + allocator.next_alloc_id = 0 + allocator.watermark_round = 0 + allocator.watermark_tail = 0 + allocator.lock = threading.Lock() + + alloc_id, offset, alloc_round = allocator.assign(60) + self.assertEqual((offset, alloc_round), (0, 0)) + allocator.free(alloc_id) + + # The old implementation left watermark=(0, 60). A 70-byte + # allocation then wrapped to round 1 and waited forever because its + # end (70) lay beyond that stale head, even though the ring was empty. + _, offset, alloc_round = allocator.assign(70) + self.assertEqual((offset, alloc_round), (0, 1)) + watermark = allocator.get_watermark() + self.assertEqual(watermark, (1, 0)) + self.assertTrue( + is_watermark_ready( + SimpleNamespace(remote_watermarks={"session": watermark}), + "session", + alloc_round, + offset + 70, + ) + ) + + def test_deferred_fanout_does_not_replay_completed_staging_destination(self): + class FiniteQueue: + def __init__(self, item): + self.items = [item] + + def get(self): + if not self.items: + raise StopIteration + return self.items.pop(0) + + def put(self, item): + self.items.append(item) + + manager = object.__new__(MooncakeKVManager) + room = 7 + chunk = TransferKVChunk( + room=room, + prefill_kv_indices=np.array([1], dtype=np.int32), + index_slice=slice(0, 1), + is_last_chunk=True, + prefill_aux_index=0, + state_indices=None, + ) + queue = FiniteQueue(chunk) + reqs = [ + SimpleNamespace( + room=room, + endpoint="127.0.0.1", + dst_port=9000 + i, + mooncake_session_id=f"session-{i}", + dst_kv_indices=np.array([1], dtype=np.int32), + dst_device_kv_indices=None, + required_dst_info_num=2, + is_dummy=False, + decode_prefix_len=0, + ) + for i in range(2) + ] + registrations = { + req.mooncake_session_id: SimpleNamespace( + requires_dcp_relayout=False, + dst_kv_ptrs=[0x1000], + dst_aux_ptrs=[0x2000], + dst_attn_tp_size=8, + dst_kv_item_len=128, + dst_kv_layer_ids=[], + staging_base_ptr=0x3000, + staging_total_size=4096, + ) + for req in reqs + } + status = {room: KVPoll.WaitingForInput} + manager.enable_trace = False + manager.enable_staging = True + manager.enable_deferred_decode_kv_release = False + manager._staging_outstanding = defaultdict(int) + manager._staging_ctx = SimpleNamespace( + prefetch_requested=set(), prefetched_rooms=set() + ) + manager.transfer_infos = {room: {req.mooncake_session_id: req for req in reqs}} + manager.req_to_decode_prefix_len = {room: 0} + manager.request_status = status + manager.check_status = lambda checked_room: status[checked_room] + manager.update_status = Mock( + side_effect=lambda checked_room, value: status.__setitem__( + checked_room, value + ) + ) + manager.sync_status_to_decode_endpoint = Mock() + manager.session_lock = threading.Lock() + manager.failed_sessions = set() + manager.session_failures = defaultdict(int) + manager.decode_kv_args_table = registrations + manager.kv_args = SimpleNamespace(kv_data_ptrs=[0x4000]) + manager.is_mla_backend = False + manager.is_hybrid_mla_backend = False + manager.attn_tp_size = 1 + manager.attn_tp_rank = 0 + manager.attn_cp_size = 1 + manager.attn_cp_rank = 0 + manager.pp_size = 1 + manager.pp_rank = 0 + manager.bootstrap_port = 7200 + manager._try_create_staging_strategy = Mock(return_value=object()) + manager._get_dsa_cache_transfer_skip_flags = Mock(return_value=(False, False)) + manager.send_aux = Mock(return_value=0) + manager.record_failure = Mock() + + transfer_calls = [] + + def transfer_side_effect(*args): + req = args[2] + transfer_calls.append(req.mooncake_session_id) + if transfer_calls == ["session-0", "session-1"]: + queue.put(chunk) + return -1, True + return 0, False + + manager._do_staging_transfer = Mock(side_effect=transfer_side_effect) + + with self.assertRaisesRegex(RuntimeError, "Transfer thread failed"): + manager.transfer_worker(queue, Mock(), staging_buffer=object()) + + self.assertEqual(transfer_calls, ["session-0", "session-1", "session-1"]) + self.assertEqual(chunk.staging_completed_sessions, {"session-0", "session-1"}) + self.assertEqual(manager.send_aux.call_count, 2) + self.assertEqual(manager.sync_status_to_decode_endpoint.call_count, 2) + self.assertEqual(status[room], KVPoll.Success) + def test_staging_response_targets_requesting_pp_rank(self): sock = Mock() receiver = SimpleNamespace(