From 8fa8d9d7e8e219e586f4a045e4649dfbac7bab04 Mon Sep 17 00:00:00 2001 From: Kevin Li <32230403+bluecoffee8@users.noreply.github.com> Date: Fri, 12 Dec 2025 04:15:00 -0800 Subject: [PATCH] [PD] Add decode PP event loop for PD disaggregation (#14945) Signed-off-by: Shangming Cai Co-authored-by: Shangming Cai --- python/sglang/srt/disaggregation/decode.py | 43 +- .../srt/disaggregation/mooncake/conn.py | 13 +- python/sglang/srt/managers/schedule_batch.py | 1 + python/sglang/srt/managers/scheduler.py | 6 +- .../sglang/srt/managers/scheduler_pp_mixin.py | 388 +++++++++++++++++- 5 files changed, 414 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ef52bda7f..b92ffce1d 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -25,7 +25,7 @@ import time from collections import deque from dataclasses import dataclass from http import HTTPStatus -from typing import TYPE_CHECKING, List, Optional, Type, Union +from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union import torch from torch.distributed import ProcessGroup @@ -199,6 +199,7 @@ class DecodePreallocQueue: bootstrap_port: int, max_total_num_tokens: int, prefill_pp_size: int, + pp_rank: int, num_reserved_decode_tokens: int, transfer_backend: TransferBackend, ): @@ -220,6 +221,7 @@ class DecodePreallocQueue: self.bootstrap_port = bootstrap_port self.max_total_num_tokens = max_total_num_tokens self.prefill_pp_size = prefill_pp_size + self.pp_rank = pp_rank self.num_reserved_decode_tokens = num_reserved_decode_tokens self.transfer_backend = transfer_backend # Queue for requests pending pre-allocation @@ -236,8 +238,7 @@ class DecodePreallocQueue: kv_args.engine_rank = self.tp_rank % (attn_tp_size) kv_args.decode_tp_size = attn_tp_size - # Note(shangming): pp is not supported on the decode side yet, so its rank is fixed to 0 - kv_args.pp_rank = 0 + kv_args.pp_rank = self.pp_rank kv_args.system_dp_rank = self.scheduler.dp_rank kv_args.prefill_pp_size = self.prefill_pp_size kv_data_ptrs, kv_data_lens, kv_item_lens = ( @@ -302,6 +303,7 @@ class DecodePreallocQueue: return if is_retracted: + req.retraction_mb_id = None self.retracted_queue.append(req) else: if req.bootstrap_host == FAKE_BOOTSTRAP_HOST: @@ -340,7 +342,9 @@ class DecodePreallocQueue: for req in reqs: self.add(req, is_retracted=is_retracted) - def resume_retracted_reqs(self) -> List[Req]: + def resume_retracted_reqs( + self, rids_to_check: Optional[List[str]] = None + ) -> List[Req]: # TODO refactor the scheduling part, reuse with the unified engine logic as much as possible # allocate memory @@ -349,6 +353,9 @@ class DecodePreallocQueue: allocatable_tokens = self._allocatable_tokens(count_retracted=False) for i, req in enumerate(self.retracted_queue): + if rids_to_check is not None and req.rid not in rids_to_check: + continue + if self.req_to_token_pool.available_size() <= 0: break @@ -377,7 +384,9 @@ class DecodePreallocQueue: return resumed_reqs - def _update_handshake_waiters(self) -> None: + def _update_handshake_waiters( + self, rids_to_check: Optional[List[str]] = None + ) -> None: if not self.queue: return @@ -389,6 +398,9 @@ class DecodePreallocQueue: ) for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + continue + if poll == KVPoll.Bootstrapping: pass elif poll == KVPoll.WaitingForInput: @@ -410,10 +422,13 @@ class DecodePreallocQueue: else: raise ValueError(f"Unexpected poll case: {poll}") - def pop_preallocated(self) -> List[DecodeRequest]: + def pop_preallocated( + self, rids_to_check: Optional[List[str]] = None + ) -> Tuple[List[DecodeRequest], List[DecodeRequest]]: """Pop the preallocated requests from the pending queue (FIFO).""" - self._update_handshake_waiters() + self._update_handshake_waiters(rids_to_check) + failed_reqs = [] preallocated_reqs = [] indices_to_remove = set() @@ -428,14 +443,20 @@ class DecodePreallocQueue: ) # First, remove all failed requests from the queue for i, decode_req in enumerate(self.queue): + if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + continue if isinstance(decode_req.req.finished_reason, FINISH_ABORT): self.scheduler.stream_output( [decode_req.req], decode_req.req.return_logprob ) + failed_reqs.append(decode_req) indices_to_remove.add(i) # Then, preallocate the remaining requests if possible for i, decode_req in enumerate(self.queue): + if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + continue + if i in indices_to_remove: continue @@ -544,7 +565,7 @@ class DecodePreallocQueue: entry for i, entry in enumerate(self.queue) if i not in indices_to_remove ] - return preallocated_reqs + return preallocated_reqs, failed_reqs @property def num_tokens_pre_allocated(self): @@ -726,7 +747,7 @@ class DecodeTransferQueue: ) decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter() - def pop_transferred(self) -> List[Req]: + def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req]: if not self.queue: return [] polls = poll_and_all_reduce( @@ -736,6 +757,8 @@ class DecodeTransferQueue: transferred_reqs = [] indices_to_remove = set() for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + continue if poll == KVPoll.Failed: error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}" try: @@ -958,7 +981,7 @@ class SchedulerDisaggregationDecodeMixin: 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() + req_conns, _ = self.disagg_decode_prealloc_queue.pop_preallocated() self.disagg_decode_transfer_queue.extend(req_conns) alloc_reqs = ( self.disagg_decode_transfer_queue.pop_transferred() diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d4414d084..56a748e52 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -807,13 +807,12 @@ class MooncakeKVManager(CommonKVManager): executor, ) - if self.pp_group.is_last_rank: - # Only the last chunk we need to send the aux data - ret = self.send_aux( - req, - kv_chunk.prefill_aux_index, - target_rank_registration_info.dst_aux_ptrs, - ) + # Only the last chunk we need to send the aux data + ret = self.send_aux( + req, + 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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c5e48128c..7d14c98da 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -677,6 +677,7 @@ class Req: # The number of times this request has been retracted / preempted. self.retraction_count = 0 + self.retraction_mb_id = None # For metrics self.metrics_collector = metrics_collector diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 606ff06d1..4a3885b30 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -901,6 +901,7 @@ class Scheduler( bootstrap_port=self.server_args.disaggregation_bootstrap_port, max_total_num_tokens=self.max_total_num_tokens, prefill_pp_size=self.server_args.disaggregation_prefill_pp, + pp_rank=self.pp_rank, num_reserved_decode_tokens=self.server_args.num_reserved_decode_tokens, transfer_backend=self.transfer_backend, ) @@ -2771,7 +2772,10 @@ def run_scheduler_process( if scheduler.enable_overlap: scheduler.event_loop_overlap_disagg_decode() else: - scheduler.event_loop_normal_disagg_decode() + if server_args.pp_size > 1: + scheduler.event_loop_pp_disagg_decode() + else: + scheduler.event_loop_normal_disagg_decode() except Exception: traceback = get_exception_traceback() diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index e1f30cead..8569c8cc0 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -512,12 +512,13 @@ class SchedulerPPMixin: # send ready PP output to rank 0 if mbs[next_first_rank_mb_id] is not None: q_event, pp_outputs_to_send = last_rank_comm_queue.popleft() - torch.cuda.current_stream().wait_event(q_event) - with torch.profiler.record_function("send_res_dict_to_next_stage"): - send_output_work = self._pp_send_dict_to_next_stage( - pp_outputs_to_send.tensors, - async_send=True, - ) + if not mbs[next_first_rank_mb_id].forward_mode.is_prebuilt(): + torch.cuda.current_stream().wait_event(q_event) + with torch.profiler.record_function("send_res_dict_to_next_stage"): + send_output_work = self._pp_send_dict_to_next_stage( + pp_outputs_to_send.tensors, + async_send=True, + ) # send the outputs from the last round to let the next stage worker run post processing if not self.pp_group.is_last_rank: if pp_outputs: @@ -549,14 +550,19 @@ class SchedulerPPMixin: if mbs[next_mb_id] is not None: with torch.profiler.record_function("recv_res_dict_from_prev_stage"): - next_pp_outputs = PPProxyTensors(self._pp_recv_dict_from_prev_stage()) - with self.copy_stream_ctx: - self.copy_stream.wait_stream(self.default_stream) - batch_result = self._pp_prep_batch_result( - mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs - ) - d2h_event = torch.cuda.Event() - d2h_event.record(torch.cuda.current_stream()) + next_pp_outputs = None + if not mbs[next_mb_id].forward_mode.is_prebuilt(): + next_pp_outputs = PPProxyTensors( + self._pp_recv_dict_from_prev_stage() + ) + if not mbs[next_mb_id].forward_mode.is_prebuilt(): + with self.copy_stream_ctx: + self.copy_stream.wait_stream(self.default_stream) + batch_result = self._pp_prep_batch_result( + mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs + ) + d2h_event = torch.cuda.Event() + d2h_event.record(torch.cuda.current_stream()) return next_pp_outputs, batch_result, d2h_event, send_output_work @@ -588,19 +594,21 @@ class SchedulerPPMixin: ) return result, event - def get_rids(self: Scheduler, req_queue: List[Req], *poll_statuses_group): + def get_rids( + self: Scheduler, req_queue: List[Req], is_send: bool, *poll_statuses_group + ): """ Used by PP, get the required rids with the given poll statuses. """ polls = poll_and_all_reduce( - [req.disagg_kv_sender for req in req_queue], + [req.disagg_kv_sender if is_send else req.kv_receiver for req in req_queue], self.tp_worker.get_attention_tp_cpu_group(), ) rids: List = [] for poll_statuses in poll_statuses_group: rids.append( [ - req.rid + req.rid if is_send else req.req.rid for req, poll in zip(req_queue, polls) if poll in poll_statuses ] @@ -760,6 +768,7 @@ class SchedulerPPMixin: # First rank, pop the bootstrap reqs from the bootstrap queue good_bootstrapped_rids, bad_bootstrapped_rids = self.get_rids( self.disagg_prefill_bootstrap_queue.queue, + True, [KVPoll.WaitingForInput], [KVPoll.Failed], ) @@ -771,6 +780,7 @@ class SchedulerPPMixin: ) curr_good_bootstrapped_rids, curr_bad_bootstrapped_rids = self.get_rids( self.disagg_prefill_bootstrap_queue.queue, + True, [KVPoll.WaitingForInput], [KVPoll.Failed], ) @@ -782,11 +792,12 @@ class SchedulerPPMixin: ) return [good_bootstrapped_rids, bad_bootstrapped_rids] - def _pp_pd_get_transferred_ids(self: Scheduler): + def _pp_pd_get_prefill_transferred_ids(self: Scheduler): # get the current stage transfer success if self.pp_group.is_first_rank: transferred_rids = self.get_rids( self.disagg_prefill_inflight_queue, + True, [KVPoll.Success, KVPoll.Failed], ) # if other ranks, do intersection with the previous rank's transferred rids @@ -797,6 +808,7 @@ class SchedulerPPMixin: # 2. get the current stage's transferred reqs info curr_transferred_rids = self.get_rids( self.disagg_prefill_inflight_queue, + True, [KVPoll.Success, KVPoll.Failed], ) # 3. new consensus rids = intersection(previous consensus rids, transfer finished rids) @@ -938,7 +950,7 @@ class SchedulerPPMixin: bmbs[mb_id] = bootstrapped_rids self._pp_commit_comm_work(send_bootstrapped_work) - transferred_rids = self._pp_pd_get_transferred_ids() + transferred_rids = self._pp_pd_get_prefill_transferred_ids() self._pp_commit_comm_work(send_transfer_work) tmbs[mb_id] = transferred_rids @@ -1052,3 +1064,341 @@ class SchedulerPPMixin: self.check_tree_cache() self.new_token_ratio = self.init_new_token_ratio self.maybe_sleep_on_idle() + + def _pp_pd_get_retract_ids(self: Scheduler, mb_id: int): + # communicate pre-consensus retracted reqs + for req in self.disagg_decode_prealloc_queue.retracted_queue: + # assign retracted reqs to the current microbatch + if req.retraction_mb_id is None: + req.retraction_mb_id = mb_id + curr_retract_rids = [ + req.rid + for req in self.disagg_decode_prealloc_queue.retracted_queue + if req.retraction_mb_id == mb_id + ] + if self.pp_group.is_first_rank: + # First rank, get all retracted req ids for the microbatch + return curr_retract_rids + else: + # Other ranks, receive the retracted reqs info from the previous rank and ensure the consensus + prev_retract_rids = self._pp_recv_pyobj_from_prev_stage() + return list(set(prev_retract_rids) & set(curr_retract_rids)) + + def _pp_pd_get_prealloc_ids(self: Scheduler): + # communicate pre-consensus prealloc reqs + if self.pp_group.is_first_rank: + # First rank, pop the preallocated reqs from the prealloc queue + good_prealloc_rids, bad_prealloc_rids = self.get_rids( + self.disagg_decode_prealloc_queue.queue, + False, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + else: + # Other ranks, receive the preallocated reqs info from the previous rank and ensure the consensus + prev_prealloc_rids = self._pp_recv_pyobj_from_prev_stage() + prev_good_prealloc_rids, prev_bad_prealloc_rids = prev_prealloc_rids + curr_good_prealloc_rids, curr_bad_prealloc_rids = self.get_rids( + self.disagg_decode_prealloc_queue.queue, + False, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + good_prealloc_rids = list( + set(prev_good_prealloc_rids) & set(curr_good_prealloc_rids) + ) + bad_prealloc_rids = list( + set(prev_bad_prealloc_rids) | set(curr_bad_prealloc_rids) + ) + return [good_prealloc_rids, bad_prealloc_rids] + + def _pp_pd_get_decode_transferred_ids(self: Scheduler): + # get the current stage transfer success + if self.pp_group.is_first_rank: + transferred_rids = self.get_rids( + self.disagg_decode_transfer_queue.queue, + False, + [KVPoll.Success, KVPoll.Failed], + ) + # if other ranks, do intersection with the previous rank's transferred rids + else: + # 2 (Release): Receive the transferred rids from the previous rank + # 1. recv previous stage's transferred reqs info + prev_transferred_rids = self._pp_recv_pyobj_from_prev_stage() + # 2. get the current stage's transferred reqs info + curr_transferred_rids = self.get_rids( + self.disagg_decode_transfer_queue.queue, + False, + [KVPoll.Success, KVPoll.Failed], + ) + # 3. new consensus rids = intersection(previous consensus rids, transfer finished rids) + transferred_rids = list( + set(prev_transferred_rids) & set(curr_transferred_rids) + ) + return transferred_rids + + def process_retract_queue(self: Scheduler, retract_rids: Optional[List[str]]): + if retract_rids is not None: + # 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( + retract_rids + ) + self.waiting_queue.extend(resumed_reqs) + return [req.rid for req in resumed_reqs] + return None + + def process_prealloc_queue(self: Scheduler, prealloc_rids: Optional[List[str]]): + if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: + # if there are still retracted requests, we do not allocate new requests + return [[], []] + + if prealloc_rids is not None: + ( + good_consensus_prealloc_rids, + bad_consensus_prealloc_rids, + ) = prealloc_rids + good_reqs, failed_reqs = self.disagg_decode_prealloc_queue.pop_preallocated( + rids_to_check=good_consensus_prealloc_rids + + bad_consensus_prealloc_rids, + ) + self.disagg_decode_transfer_queue.extend(good_reqs) + return [ + [req.req.rid for req in good_reqs], + [req.req.rid for req in failed_reqs], + ] + return None + + def process_decode_transfer_queue( + self: Scheduler, release_rids: Optional[List[str]] + ): + if release_rids is not None: + released_reqs = self.disagg_decode_transfer_queue.pop_transferred( + release_rids + ) + self.waiting_queue.extend(released_reqs) + return [req.rid for req in released_reqs] + return None + + @DynamicGradMode() + def event_loop_pp_disagg_decode(self: Scheduler): + self.pp_loop_size: int = self.pp_size + self.server_args.pp_async_batch_depth + mbs = [None] * self.pp_loop_size + last_mbs = [None] * self.pp_loop_size + self.running_mbs = [ + ScheduleBatch(reqs=[], batch_is_full=False) + for _ in range(self.pp_loop_size) + ] + mb_metadata: List[Optional[PPBatchMetadata]] = [None] * self.pp_loop_size + pp_outputs: Optional[PPProxyTensors] = None + last_rank_comm_queue: deque[Tuple[torch.cuda.Event, PPProxyTensors]] = deque() + + # PD additional + + # consensus rids + consensus_retract_rids: Optional[List[str]] = None + consensus_prealloc_rids: Optional[List[str]] = None + release_rids: Optional[List[str]] = None # consensus transferred rids + + rmbs = [None] * self.pp_loop_size + pmbs = [None] * self.pp_loop_size + tmbs = [None] * self.pp_loop_size + + send_req_work = [] + + # send info to reach consensus + send_retract_work = [] + send_prealloc_work = [] + send_transfer_work = [] + + # send consensus info + send_consensus_retract_work = [] + send_consensus_prealloc_work = [] + send_release_work = [] + + send_proxy_work = [] + send_output_work = [] + + while True: + server_is_idle = True + for mb_id in range(self.pp_loop_size): + self.running_batch = self.running_mbs[mb_id] + self.last_batch = last_mbs[mb_id] + next_first_rank_mb_id = (mb_id + self.pp_size) % self.pp_loop_size + next_mb_id = (mb_id + 1) % self.pp_loop_size + + next_pp_outputs = None + next_consensus_retract_rids = None + next_consensus_prealloc_rids = None + next_release_rids = None + d2h_event = None + next_batch_result = None + + recv_reqs = self.recv_requests() + self.process_input_requests(recv_reqs) + + if not self.pp_group.is_last_rank: + self._pp_commit_comm_work(send_req_work) + + # reaching consensus through PP ranks + retract_rids = self._pp_pd_get_retract_ids(mb_id) + rmbs[mb_id] = retract_rids + self._pp_commit_comm_work(send_retract_work) + + prealloc_rids = self._pp_pd_get_prealloc_ids() + pmbs[mb_id] = prealloc_rids + self._pp_commit_comm_work(send_prealloc_work) + + transferred_rids = self._pp_pd_get_decode_transferred_ids() + tmbs[mb_id] = transferred_rids + self._pp_commit_comm_work(send_transfer_work) + + # get batch to run and proxy tensors if needed + batch = self.get_next_disagg_decode_batch_to_run() + mbs[mb_id] = batch + self.running_mbs[mb_id] = self.running_batch + + self.cur_batch: Optional[ScheduleBatch] = mbs[mb_id] + if self.cur_batch: + server_is_idle = False + pp_proxy_tensors = None + if not self.cur_batch.forward_mode.is_prebuilt(): + pp_proxy_tensors = self._pp_recv_proxy_tensors() + + # early send output if possible + if self.server_args.pp_async_batch_depth > 0: + self._pp_commit_comm_work(work=send_output_work) + next_pp_outputs, next_batch_result, d2h_event, send_output_work = ( + self._pp_send_recv_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + mbs, + mb_metadata, + last_rank_comm_queue, + pp_outputs, + ) + ) + self._pp_commit_comm_work(send_proxy_work) + # run batch + if self.cur_batch: + result, event = self._pp_launch_batch( + mb_id, pp_proxy_tensors, mb_metadata, last_rank_comm_queue + ) + # regular send output + if self.server_args.pp_async_batch_depth == 0: + self._pp_commit_comm_work(work=send_output_work) + next_pp_outputs, next_batch_result, d2h_event, send_output_work = ( + self._pp_send_recv_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + mbs, + mb_metadata, + last_rank_comm_queue, + pp_outputs, + ) + ) + + # reach consensus on last rank and send to PP=0 + # otherwise, just pass along previous consensus + send_consensus_retract_work, consensus_retract_rids = ( + self._pp_pd_send_consensus_bootstrapped_ids( # reuse the function + rmbs, + next_first_rank_mb_id, + consensus_retract_rids, + retract_rids, + ) + ) + + send_consensus_prealloc_work, consensus_prealloc_rids = ( + self._pp_pd_send_consensus_bootstrapped_ids( # reuse the function + pmbs, + next_first_rank_mb_id, + consensus_prealloc_rids, + prealloc_rids, + ) + ) + + send_release_work, release_rids = ( + self._pp_pd_send_consensus_release_ids( + tmbs, next_first_rank_mb_id, release_rids, transferred_rids + ) + ) + + if self.server_args.disaggregation_decode_enable_offload_kvcache: + self.decode_offload_manager.check_offload_progress() + + if rmbs[next_mb_id] is not None: + next_consensus_retract_rids = self._pp_recv_pyobj_from_prev_stage() + next_consensus_retract_rids = self.process_retract_queue( + next_consensus_retract_rids + ) + self._pp_commit_comm_work(send_consensus_retract_work) + + if pmbs[next_mb_id] is not None: + next_consensus_prealloc_rids = self._pp_recv_pyobj_from_prev_stage() + next_consensus_prealloc_rids = self.process_prealloc_queue( + next_consensus_prealloc_rids + ) + self._pp_commit_comm_work(send_consensus_prealloc_work) + + if tmbs[next_mb_id] is not None: + next_release_rids = self._pp_recv_pyobj_from_prev_stage() + next_release_rids = self.process_decode_transfer_queue( + next_release_rids + ) + self._pp_commit_comm_work(send_release_work) + + # post-process the coming microbatch + if mbs[next_mb_id] is not None: + if not mbs[next_mb_id].forward_mode.is_prebuilt(): + d2h_event.synchronize() + self._pp_process_batch_result( + mbs[next_mb_id], + next_batch_result, + ) + last_mbs[next_mb_id] = mbs[next_mb_id] + + if not self.pp_group.is_last_rank: + send_req_work = self._pp_send_pyobj_to_next_stage( + recv_reqs, async_send=True + ) + send_retract_work = self._pp_send_pyobj_to_next_stage( + retract_rids, async_send=True + ) + send_prealloc_work = self._pp_send_pyobj_to_next_stage( + prealloc_rids, async_send=True + ) + send_transfer_work = self._pp_send_pyobj_to_next_stage( + transferred_rids, async_send=True + ) + if self.cur_batch and not self.cur_batch.forward_mode.is_prebuilt(): + torch.cuda.current_stream().wait_event(event) + send_proxy_work = self._pp_send_dict_to_next_stage( + result.pp_hidden_states_proxy_tensors.tensors, + async_send=True, + ) + + if hasattr(self, "delayed_weight_sync_fn"): + self.delayed_weight_sync_fn() + self.delayed_weight_sync_fn = None + + pp_outputs = next_pp_outputs + release_rids = next_release_rids + consensus_retract_rids = next_consensus_retract_rids + consensus_prealloc_rids = next_consensus_prealloc_rids + + self.running_batch.batch_is_full = False + + # When the server is idle, self-check and re-init some states + queue_size = ( + len(self.waiting_queue) + + len(self.disagg_decode_transfer_queue.queue) + + len(self.disagg_decode_prealloc_queue.queue) + ) + if self.server_args.disaggregation_decode_enable_offload_kvcache: + queue_size += len(self.decode_offload_manager.ongoing_offload) + + if server_is_idle and queue_size == 0: + self.check_memory() + self.check_tree_cache() + self.new_token_ratio = self.init_new_token_ratio + self.maybe_sleep_on_idle()