From eeb2b9b2596cad7ca335fbdf669332d6b56d5e0f Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Wed, 17 Dec 2025 23:31:25 +0800 Subject: [PATCH] [PP] Minor code cleanup for Pipeline Parallelism (#15329) Signed-off-by: Shangming Cai --- .../sglang/srt/managers/scheduler_pp_mixin.py | 1633 ++++++++--------- 1 file changed, 795 insertions(+), 838 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 8569c8cc0..dc13fb840 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -31,184 +31,496 @@ if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler -class ChunkSizePredictor: - """ - Predictor for dynamic chunk size based on quadratic latency model. - - Models latency as: f(l) = a*l^2 + b*l + c - Predicts next chunk size x such that: f(L+x) - f(L) = target_latency - """ - - def __init__(self): - self.quadratic_coeff_a = 0.0 - self.linear_coeff_b = 0.0 - self.constant_coeff_c = 0.0 - self.target_latency: Optional[float] = None - self.is_ready = False - - def fit(self, seq_lens: List[int], latencies: List[float]): - """Fit quadratic coefficients f(l) = al^2 + bl + c from data points.""" - L = np.array(seq_lens, dtype=np.float64) - T = np.array(latencies, dtype=np.float64) - - if len(L) < 8: - raise ValueError( - f"Not enough data points for quadratic fitting ({len(L)} < 8). " - "Need at least 8 samples with different sequence lengths." - ) - - # Build design matrix for f(l) = al^2 + bl + c - X = np.column_stack([L * L, L, np.ones_like(L)]) # [l^2, l, 1] - - try: - coeffs, residuals, rank, s = np.linalg.lstsq(X, T, rcond=None) - if len(coeffs) >= 3: - fitted_a = float(coeffs[0]) # quadratic coefficient - fitted_b = float(coeffs[1]) # linear coefficient - fitted_c = float(coeffs[2]) # constant coefficient - else: - raise ValueError("Failed to fit coefficients: insufficient rank") - except np.linalg.LinAlgError as e: - raise ValueError(f"Failed to fit f(l) = al^2 + bl + c: {e}") - - # Validate coefficients - if fitted_a <= 0: - raise ValueError( - f"Fitted quadratic coefficient a={fitted_a:.2e} is not positive. " - "Attention has O(n^2) complexity, so a must be positive. " - "Check warmup data quality." - ) - - if fitted_b < 0: - logger.warning( - f"Fitted linear coefficient b={fitted_b:.2e} is negative. Setting b=0." - ) - fitted_b = 0.0 - - self.quadratic_coeff_a = fitted_a - self.linear_coeff_b = fitted_b - self.constant_coeff_c = fitted_c - - logger.info( - f"[ChunkSizePredictor] Fitted coefficients: a={fitted_a:.2e}, " - f"b={fitted_b:.2e}, c={fitted_c:.2e}" - ) - - def set_target_latency(self, base_chunk_size: int): - """Set target latency based on base chunk size: target = f(base_chunk_size) - f(0).""" - - def f(l: float) -> float: - """Total latency function: f(l) = al^2 + bl + c (or bl + c for linear)""" - return ( - self.quadratic_coeff_a * l * l - + self.linear_coeff_b * l - + self.constant_coeff_c - ) - - self.target_latency = f(float(base_chunk_size)) - f(0.0) - - if self.target_latency <= 0: - raise ValueError( - f"Calculated target_latency={self.target_latency:.2f}ms is not positive. " - "Check warmup data quality." - ) - - logger.info( - f"[ChunkSizePredictor] Target latency: {self.target_latency:.2f}ms " - f"(base_chunk_size={base_chunk_size})" - ) - - def predict_next_chunk_size( - self, - history_len: int, - base_chunk_size: int, - page_size: int, - context_len: int, - max_chunk_size: Optional[int] = None, - ) -> Optional[int]: - """ - Predict next chunk size x such that f(history_len + x) - f(history_len) = target_latency. - - Args: - history_len: Current sequence length (L) - base_chunk_size: Base chunk size - page_size: Page size for alignment - context_len: Maximum context length - max_chunk_size: Maximum allowed chunk size (optional) - - Returns: - Predicted chunk size, or None if prediction fails - """ - if not self.is_ready or self.target_latency is None: - return None - - # Handle quadratic model: f(l) = al^2 + bl + c - if self.quadratic_coeff_a <= 0: - return None - - # Solve f(L+x) - f(L) = T - # where f(L) = a*L^2 + b*L + c - # This expands to: ax^2 + (2aL+b)x - T = 0 - # A = a, B = 2aL + b, C = -T - A = self.quadratic_coeff_a - B = 2 * self.quadratic_coeff_a * history_len + self.linear_coeff_b - C = -self.target_latency - - discriminant = B * B - 4 * A * C - - if discriminant < 0: - logger.warning( - f"Discriminant is negative ({discriminant:.2e}). " - f"No real solution for chunk size. L={history_len}, T={self.target_latency:.2f}ms." - ) - return None - - sqrt_discriminant = math.sqrt(discriminant) - calculated_chunk_size_float = (-B + sqrt_discriminant) / (2 * A) - - if calculated_chunk_size_float <= 0: - logger.warning( - f"Calculated chunk size is non-positive ({calculated_chunk_size_float:.2f}). " - f"L={history_len}, T={self.target_latency:.2f}ms." - ) - return None - - # Use a smooth coefficient to reduce the abrupt decrease in chunk size - smooth_coeff = envs.SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR.get() - smoothed_chunk_size = base_chunk_size + smooth_coeff * ( - calculated_chunk_size_float - base_chunk_size - ) - calculated_chunk_size = int(smoothed_chunk_size) - - # Align to page_size (round down to nearest multiple) - alignment_size = max(page_size, 1) - dynamic_chunk_size = (calculated_chunk_size // alignment_size) * alignment_size - - # Ensure aligned size is at least alignment_size - if dynamic_chunk_size < alignment_size: - dynamic_chunk_size = alignment_size - - # Apply constraints - max_allowed = context_len - history_len - 100 # Leave 100 tokens margin - if max_chunk_size is not None: - max_allowed = min(max_allowed, max_chunk_size) - dynamic_chunk_size = min(dynamic_chunk_size, max_allowed) - - # Align again after min operation - dynamic_chunk_size = (dynamic_chunk_size // alignment_size) * alignment_size - - if dynamic_chunk_size < alignment_size: - return None - - return dynamic_chunk_size - - @dataclass class PPBatchMetadata: can_run_cuda_graph: bool class SchedulerPPMixin: + @DynamicGradMode() + def event_loop_pp(self: Scheduler): + """ + A scheduler loop for pipeline parallelism. + Notes: + 1. Each stage runs in the same order and is notified by the previous stage. + 2. We use async send but sync recv to avoid desynchronization while minimizing the communication overhead. + 3. We can use async batch depth to buffer the outputs in the last stage for to allow overlapping the GPU computation and CPU processing and avoid last PP rank staggler. + + Unified Schedule: + ==================================================================== + Stage P + recv ith req from previous stage + recv ith proxy from previous stage + run ith batch + recv prev (i+1)% mb_size th outputs + process batch result of prev (i+1)% mb_size th batch (can be run in parallel with the curr batch GPU computation) + send ith req to next stage + send ith proxy to next stage + send current stage's outputs to next stage(can be stashed and delayed to send later) + + the above order can be optimized and reordered to minimize communication-related CPU stall and overhead bubbles. + + ==================================================================== + """ + self.init_pp_loop_state() + 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 = self.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 + with torch.profiler.record_function("recv_requests"): + recv_reqs = self.recv_requests() + self.process_input_requests(recv_reqs) + if not self.pp_group.is_last_rank: + self._pp_commit_comm_work(self.send_req_work) + with torch.profiler.record_function("send_reqs_to_next_stage"): + self.send_req_work = self._pp_send_pyobj_to_next_stage( + recv_reqs, + async_send=True, + ) + with torch.profiler.record_function("get_next_batch_to_run"): + self.mbs[mb_id] = self.get_next_batch_to_run() + self.running_mbs[mb_id] = self.running_batch + self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + if self.cur_batch: + server_is_idle = False + pp_proxy_tensors = self._pp_recv_proxy_tensors() + next_pp_outputs = None + next_batch_result = None + d2h_event = None + if self.server_args.pp_async_batch_depth > 0: + next_pp_outputs, next_batch_result, d2h_event = ( + self._pp_commit_send_output_work_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + ) + ) + self._pp_commit_comm_work(self.send_proxy_work) + if self.cur_batch: + result, self.launch_event = self._pp_launch_batch( + mb_id, + pp_proxy_tensors, + self.mb_metadata, + self.last_rank_comm_queue, + ) + if self.server_args.pp_async_batch_depth == 0: + next_pp_outputs, next_batch_result, d2h_event = ( + self._pp_commit_send_output_work_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + ) + ) + if self.mbs[next_mb_id] is not None: + d2h_event.synchronize() + with torch.profiler.record_function("process_batch_result"): + self._pp_process_batch_result( + self.mbs[next_mb_id], + next_batch_result, + ) + self.last_mbs[next_mb_id] = self.mbs[next_mb_id] + if not self.pp_group.is_last_rank: + if self.cur_batch: + torch.cuda.current_stream().wait_event(self.launch_event) + with torch.profiler.record_function( + "send_proxy_dict_to_next_stage" + ): + self.send_proxy_work = self._pp_send_dict_to_next_stage( + result.pp_hidden_states_proxy_tensors.tensors, + async_send=True, + ) + + self.pp_outputs = next_pp_outputs + + # When the server is idle, self-check and re-init some states + if server_is_idle: + self.check_during_pp_idle() + + @DynamicGradMode() + def event_loop_pp_disagg_prefill(self: Scheduler): + """ + This is the prefill server event loop for pipeline parallelism. + + Notes: + 1. Following the same rules as the event_loop_pp. + 2. Adds extra steps for KV transfer process: bootstrap + release. + + Prefill Server Schedule: + ==================================================================== + Stage P + recv ith req from previous stage + recv ith bootstrap req from previous stage + recv ith transferred req from previous stage + recv ith proxy from previous stage + run ith batch + recv prev (i+1) % mb_size th consensus bootstrapped req from previous stage + local consensus on bootstrapped req + recv prev (i+1) % mb_size th release req from previous stage + local consensus on release req + recv prev (i+1) % mb_size th outputs + process batch result of prev (i+1)% mb_size th batch (can be run in parallel with the curr batch GPU computation) + send ith req to next stage + send ith bootstrap req to next stage + send ith transferred req to next stage + send ith proxy to next stage + send current stage's outputs to next stage (can be stashed and delayed to send later) + + the above order can be optimized and reordered to minimize communication-related CPU stall and overhead bubbles. + ==================================================================== + + There are two additional elements compared to the regular schedule: + + Bootstrap Requests + Release Requests: + - Both can have local failure and need to be consensus on. PP needs to guarantee eventual consistency of local failure and flush malfunc requests out as soft error. + + """ + self.init_pp_loop_state() + + # PD additional state initialization + bmbs = [None] * self.pp_loop_size + tmbs = [None] * self.pp_loop_size + consensus_bootstrapped_rids: Optional[List[str]] = None + transferred_rids: List[str] = [] + release_rids: Optional[List[str]] = None + send_bootstrapped_work = [] + send_transfer_work = [] + send_consensus_bootstrapped_work = [] + send_release_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 = self.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_release_rids = None + next_consensus_bootstrapped_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(self.send_req_work) + + bootstrapped_rids = self._pp_pd_get_bootstrapped_ids() + bmbs[mb_id] = bootstrapped_rids + self._pp_commit_comm_work(send_bootstrapped_work) + + transferred_rids = self._pp_pd_get_prefill_transferred_ids() + self._pp_commit_comm_work(send_transfer_work) + tmbs[mb_id] = transferred_rids + + self.process_prefill_chunk() + batch = self.get_new_batch_prefill() + if self.require_mlp_sync: + batch = self.prepare_mlp_sync_batch(batch) + self.mbs[mb_id] = batch + self.running_mbs[mb_id] = self.running_batch + + self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + if self.cur_batch: + server_is_idle = False + pp_proxy_tensors = self._pp_recv_proxy_tensors() + + if self.server_args.pp_async_batch_depth > 0: + next_pp_outputs, next_batch_result, d2h_event = ( + self._pp_commit_send_output_work_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + ) + ) + self._pp_commit_comm_work(self.send_proxy_work) + if self.cur_batch: + result, self.launch_event = self._pp_launch_batch( + mb_id, + pp_proxy_tensors, + self.mb_metadata, + self.last_rank_comm_queue, + ) + if self.server_args.pp_async_batch_depth == 0: + next_pp_outputs, next_batch_result, d2h_event = ( + self._pp_commit_send_output_work_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + ) + ) + send_consensus_bootstrapped_work, consensus_bootstrapped_rids = ( + self._pp_pd_send_consensus_bootstrapped_ids( + bmbs, + next_first_rank_mb_id, + consensus_bootstrapped_rids, + bootstrapped_rids, + ) + ) + send_release_work, release_rids = ( + self._pp_pd_send_consensus_release_ids( + tmbs, next_first_rank_mb_id, release_rids, transferred_rids + ) + ) + + if bmbs[next_mb_id] is not None: + next_consensus_bootstrapped_rids = ( + self._pp_recv_pyobj_from_prev_stage() + ) + next_consensus_bootstrapped_rids = self.process_bootstrapped_queue( + next_consensus_bootstrapped_rids + ) + self._pp_commit_comm_work(send_consensus_bootstrapped_work) + if tmbs[next_mb_id] is not None: + next_release_rids = self._pp_recv_pyobj_from_prev_stage() + self._pp_commit_comm_work(send_release_work) + # post-process the coming microbatch + if self.mbs[next_mb_id] is not None: + d2h_event.synchronize() + self._pp_process_batch_result( + self.mbs[next_mb_id], + next_batch_result, + ) + self.last_mbs[next_mb_id] = self.mbs[next_mb_id] + + if tmbs[next_mb_id] is not None: + self.process_disagg_prefill_inflight_queue(next_release_rids) + if not self.pp_group.is_last_rank: + self.send_req_work = self._pp_send_pyobj_to_next_stage( + recv_reqs, async_send=True + ) + send_bootstrapped_work = self._pp_send_pyobj_to_next_stage( + bootstrapped_rids, async_send=True + ) + send_transfer_work = self._pp_send_pyobj_to_next_stage( + transferred_rids, async_send=True + ) + if self.cur_batch: + torch.cuda.current_stream().wait_event(self.launch_event) + self.send_proxy_work = self._pp_send_dict_to_next_stage( + result.pp_hidden_states_proxy_tensors.tensors, + async_send=True, + ) + + self.pp_outputs = next_pp_outputs + release_rids = next_release_rids + consensus_bootstrapped_rids = next_consensus_bootstrapped_rids + + self.running_batch.batch_is_full = False + + # When the server is idle, self-check and re-init some states + if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0: + self.check_during_pp_idle() + + @DynamicGradMode() + def event_loop_pp_disagg_decode(self: Scheduler): + self.init_pp_loop_state() + + # PD additional state initialization + rmbs = [None] * self.pp_loop_size + pmbs = [None] * self.pp_loop_size + tmbs = [None] * self.pp_loop_size + consensus_retract_rids: Optional[List[str]] = None + consensus_prealloc_rids: Optional[List[str]] = None + release_rids: Optional[List[str]] = None # consensus transferred rids + send_retract_work = [] + send_prealloc_work = [] + send_transfer_work = [] + send_consensus_retract_work = [] + send_consensus_prealloc_work = [] + send_release_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 = self.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(self.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() + self.mbs[mb_id] = batch + self.running_mbs[mb_id] = self.running_batch + + self.cur_batch: Optional[ScheduleBatch] = self.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: + next_pp_outputs, next_batch_result, d2h_event = ( + self._pp_commit_send_output_work_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + ) + ) + self._pp_commit_comm_work(self.send_proxy_work) + + if self.cur_batch: + result, self.launch_event = self._pp_launch_batch( + mb_id, + pp_proxy_tensors, + self.mb_metadata, + self.last_rank_comm_queue, + ) + + if self.server_args.pp_async_batch_depth == 0: + next_pp_outputs, next_batch_result, d2h_event = ( + self._pp_commit_send_output_work_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + ) + ) + + # 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( + 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( + 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 self.mbs[next_mb_id] is not None: + if not self.mbs[next_mb_id].forward_mode.is_prebuilt(): + d2h_event.synchronize() + self._pp_process_batch_result( + self.mbs[next_mb_id], + next_batch_result, + ) + self.last_mbs[next_mb_id] = self.mbs[next_mb_id] + + if not self.pp_group.is_last_rank: + self.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(self.launch_event) + self.send_proxy_work = self._pp_send_dict_to_next_stage( + result.pp_hidden_states_proxy_tensors.tensors, + async_send=True, + ) + + self.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_during_pp_idle() + + def init_pp_loop_state(self: Scheduler): + self.pp_loop_size: int = self.pp_size + self.server_args.pp_async_batch_depth + self.mbs = [None] * self.pp_loop_size + self.last_mbs = [None] * self.pp_loop_size + self.running_mbs = [ + ScheduleBatch(reqs=[], batch_is_full=False) + for _ in range(self.pp_loop_size) + ] + self.mb_metadata: List[Optional[PPBatchMetadata]] = [None] * self.pp_loop_size + self.pp_outputs: Optional[PPProxyTensors] = None + self.last_rank_comm_queue: deque[Tuple[torch.cuda.Event, PPProxyTensors]] = ( + deque() + ) + + self.send_req_work = [] + self.send_proxy_work = [] + self.send_output_work = [] + self.launch_event = None + def profile_and_init_predictor(self: Scheduler): """ Profile prefill latency for dynamic chunk sizing. @@ -375,11 +687,158 @@ class SchedulerPPMixin: return predicted_size + def check_during_pp_idle(self: Scheduler): + self.check_memory() + self.check_tree_cache() + self.new_token_ratio = self.init_new_token_ratio + self.maybe_sleep_on_idle() + + def process_bootstrapped_queue( + self: Scheduler, bootstrapped_rids: Optional[List[str]] + ): + # finished consensus bootstrapped reqs and prepare the waiting queue + if bootstrapped_rids is not None: + ( + good_consensus_bootstrapped_rids, + bad_consensus_bootstrapped_rids, + ) = bootstrapped_rids + good_reqs, failed_reqs = ( + self.disagg_prefill_bootstrap_queue.pop_bootstrapped( + return_failed_reqs=True, + rids_to_check=good_consensus_bootstrapped_rids + + bad_consensus_bootstrapped_rids, + ) + ) + self.waiting_queue.extend(good_reqs) + return [[req.rid for req in good_reqs], [req.rid for req in failed_reqs]] + return None + + def _pp_pd_get_bootstrapped_ids(self: Scheduler): + # communicate pre-consensus bootstrapp reqs + if self.pp_group.is_first_rank: + # 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], + ) + else: + # Other ranks, receive the bootstrap reqs info from the previous rank and ensure the consensus + prev_bootstrapped_rids = self._pp_recv_pyobj_from_prev_stage() + prev_good_bootstrapped_rids, prev_bad_bootstrapped_rids = ( + prev_bootstrapped_rids + ) + curr_good_bootstrapped_rids, curr_bad_bootstrapped_rids = self.get_rids( + self.disagg_prefill_bootstrap_queue.queue, + True, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + good_bootstrapped_rids = list( + set(prev_good_bootstrapped_rids) & set(curr_good_bootstrapped_rids) + ) + bad_bootstrapped_rids = list( + set(prev_bad_bootstrapped_rids) | set(curr_bad_bootstrapped_rids) + ) + return [good_bootstrapped_rids, bad_bootstrapped_rids] + + 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 + 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_prefill_inflight_queue, + True, + [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 _pp_pd_send_consensus_bootstrapped_ids( + self: Scheduler, + bmbs: List[List[str]], + next_first_rank_mb_id: int, + consensus_bootstrapped_rids: List[str], + bootstrapped_rids: List[str], + ): + # 3 (Release): send the release rids from last stage to the first stage + send_consensus_bootstrapped_work = [] + if self.pp_group.is_last_rank: + if bmbs[next_first_rank_mb_id] is not None: + consensus_bootstrapped_rids = bootstrapped_rids + send_consensus_bootstrapped_work = self._pp_send_pyobj_to_next_stage( + consensus_bootstrapped_rids, async_send=True + ) + # 4 (Release): send the release rids from non last rank to the next rank + else: + if consensus_bootstrapped_rids is not None: + send_consensus_bootstrapped_work = self._pp_send_pyobj_to_next_stage( + consensus_bootstrapped_rids, async_send=True + ) + return send_consensus_bootstrapped_work, consensus_bootstrapped_rids + + def _pp_pd_send_consensus_release_ids( + self: Scheduler, + tmbs: List[List[str]], + next_first_rank_mb_id: int, + release_rids: List[str], + transferred_rids: List[str], + ): + send_release_work = [] + if self.pp_group.is_last_rank: + if tmbs[next_first_rank_mb_id] is not None: + release_rids = transferred_rids + send_release_work = self._pp_send_pyobj_to_next_stage( + release_rids, async_send=True + ) + # 4 (Release): send the release rids from non last rank to the next rank + else: + if release_rids is not None: + send_release_work = self._pp_send_pyobj_to_next_stage( + release_rids, async_send=True + ) + return send_release_work, release_rids + def _pp_commit_comm_work(self: Scheduler, work: List[P2PWork]) -> None: for p2p_work in work: p2p_work.work.wait() work.clear() + def _pp_commit_send_output_work_and_preprocess_output_tensors( + self: Scheduler, + next_first_rank_mb_id: int, + next_mb_id: int, + ) -> Tuple[PPProxyTensors, GenerationBatchResult, torch.cuda.Event]: + self._pp_commit_comm_work(work=self.send_output_work) + ( + next_pp_outputs, + next_batch_result, + d2h_event, + self.send_output_work, + ) = self._pp_send_recv_and_preprocess_output_tensors( + next_first_rank_mb_id, + next_mb_id, + self.mbs, + self.mb_metadata, + self.last_rank_comm_queue, + self.pp_outputs, + ) + return next_pp_outputs, next_batch_result, d2h_event + def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False): p2p_work = [] if self.attn_tp_rank == 0: @@ -615,456 +1074,6 @@ class SchedulerPPMixin: ) return tuple(rids) if len(rids) > 1 else rids[0] - @DynamicGradMode() - def event_loop_pp(self: Scheduler): - """ - A scheduler loop for pipeline parallelism. - Notes: - 1. Each stage runs in the same order and is notified by the previous stage. - 2. We use async send but sync recv to avoid desynchronization while minimizing the communication overhead. - 3. We can use async batch depth to buffer the outputs in the last stage for to allow overlapping the GPU computation and CPU processing and avoid last PP rank staggler. - - Unified Schedule: - ==================================================================== - Stage P - recv ith req from previous stage - recv ith proxy from previous stage - run ith batch - recv prev (i+1)% mb_size th outputs - process batch result of prev (i+1)% mb_size th batch (can be run in parallel with the curr batch GPU computation) - send ith req to next stage - send ith proxy to next stage - send current stage's outputs to next stage(can be stashed and delayed to send later) - - the above order can be optimized and reordered to minimize communication-related CPU stall and overhead bubbles. - - ==================================================================== - """ - 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() - send_req_work = [] - send_proxy_work = [] - send_output_work = [] - event = None - 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 - with torch.profiler.record_function("recv_requests"): - 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) - with torch.profiler.record_function("send_reqs_to_next_stage"): - send_req_work = self._pp_send_pyobj_to_next_stage( - recv_reqs, - async_send=True, - ) - with torch.profiler.record_function("get_next_batch_to_run"): - mbs[mb_id] = self.get_next_batch_to_run() - 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 = self._pp_recv_proxy_tensors() - next_pp_outputs = None - next_batch_result = None - d2h_event = None - 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) - if self.cur_batch: - result, event = self._pp_launch_batch( - mb_id, pp_proxy_tensors, mb_metadata, last_rank_comm_queue - ) - 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, - ) - ) - if mbs[next_mb_id] is not None: - d2h_event.synchronize() - with torch.profiler.record_function("process_batch_result"): - 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: - if self.cur_batch: - torch.cuda.current_stream().wait_event(event) - with torch.profiler.record_function( - "send_proxy_dict_to_next_stage" - ): - send_proxy_work = self._pp_send_dict_to_next_stage( - result.pp_hidden_states_proxy_tensors.tensors, - async_send=True, - ) - - # if self.delayed_weight_sync_fn: - # self.delayed_weight_sync_fn() - # self.delayed_weight_sync_fn = None - - pp_outputs = next_pp_outputs - - # When the server is idle, self-check and re-init some states - if server_is_idle: - self.check_memory() - self.check_tree_cache() - self.new_token_ratio = self.init_new_token_ratio - self.maybe_sleep_on_idle() - - def process_bootstrapped_queue( - self: Scheduler, bootstrapped_rids: Optional[List[str]] - ): - # finished consensus bootstrapped reqs and prepare the waiting queue - if bootstrapped_rids is not None: - ( - good_consensus_bootstrapped_rids, - bad_consensus_bootstrapped_rids, - ) = bootstrapped_rids - good_reqs, failed_reqs = ( - self.disagg_prefill_bootstrap_queue.pop_bootstrapped( - return_failed_reqs=True, - rids_to_check=good_consensus_bootstrapped_rids - + bad_consensus_bootstrapped_rids, - ) - ) - self.waiting_queue.extend(good_reqs) - return [[req.rid for req in good_reqs], [req.rid for req in failed_reqs]] - return None - - def _pp_pd_get_bootstrapped_ids(self: Scheduler): - # communicate pre-consensus bootstrapp reqs - if self.pp_group.is_first_rank: - # 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], - ) - else: - # Other ranks, receive the bootstrap reqs info from the previous rank and ensure the consensus - prev_bootstrapped_rids = self._pp_recv_pyobj_from_prev_stage() - prev_good_bootstrapped_rids, prev_bad_bootstrapped_rids = ( - prev_bootstrapped_rids - ) - curr_good_bootstrapped_rids, curr_bad_bootstrapped_rids = self.get_rids( - self.disagg_prefill_bootstrap_queue.queue, - True, - [KVPoll.WaitingForInput], - [KVPoll.Failed], - ) - good_bootstrapped_rids = list( - set(prev_good_bootstrapped_rids) & set(curr_good_bootstrapped_rids) - ) - bad_bootstrapped_rids = list( - set(prev_bad_bootstrapped_rids) | set(curr_bad_bootstrapped_rids) - ) - return [good_bootstrapped_rids, bad_bootstrapped_rids] - - 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 - 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_prefill_inflight_queue, - True, - [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 _pp_pd_send_consensus_bootstrapped_ids( - self: Scheduler, - bmbs: List[List[str]], - next_first_rank_mb_id: int, - consensus_bootstrapped_rids: List[str], - bootstrapped_rids: List[str], - ): - # 3 (Release): send the release rids from last stage to the first stage - send_consensus_bootstrapped_work = [] - if self.pp_group.is_last_rank: - if bmbs[next_first_rank_mb_id] is not None: - consensus_bootstrapped_rids = bootstrapped_rids - send_consensus_bootstrapped_work = self._pp_send_pyobj_to_next_stage( - consensus_bootstrapped_rids, async_send=True - ) - # 4 (Release): send the release rids from non last rank to the next rank - else: - if consensus_bootstrapped_rids is not None: - send_consensus_bootstrapped_work = self._pp_send_pyobj_to_next_stage( - consensus_bootstrapped_rids, async_send=True - ) - return send_consensus_bootstrapped_work, consensus_bootstrapped_rids - - def _pp_pd_send_consensus_release_ids( - self: Scheduler, - tmbs: List[List[str]], - next_first_rank_mb_id: int, - release_rids: List[str], - transferred_rids: List[str], - ): - send_release_work = [] - if self.pp_group.is_last_rank: - if tmbs[next_first_rank_mb_id] is not None: - release_rids = transferred_rids - send_release_work = self._pp_send_pyobj_to_next_stage( - release_rids, async_send=True - ) - # 4 (Release): send the release rids from non last rank to the next rank - else: - if release_rids is not None: - send_release_work = self._pp_send_pyobj_to_next_stage( - release_rids, async_send=True - ) - return send_release_work, release_rids - - @DynamicGradMode() - def event_loop_pp_disagg_prefill(self: Scheduler): - """ - This is the prefill server event loop for pipeline parallelism. - - Notes: - 1. Following the same rules as the event_loop_pp. - 2. Adds extra steps for KV transfer process: bootstrap + release. - - Prefill Server Schedule: - ==================================================================== - Stage P - recv ith req from previous stage - recv ith bootstrap req from previous stage - recv ith transferred req from previous stage - recv ith proxy from previous stage - run ith batch - recv prev (i+1) % mb_size th consensus bootstrapped req from previous stage - local consensus on bootstrapped req - recv prev (i+1) % mb_size th release req from previous stage - local consensus on release req - recv prev (i+1) % mb_size th outputs - process batch result of prev (i+1)% mb_size th batch (can be run in parallel with the curr batch GPU computation) - send ith req to next stage - send ith bootstrap req to next stage - send ith transferred req to next stage - send ith proxy to next stage - send current stage's outputs to next stage (can be stashed and delayed to send later) - - the above order can be optimized and reordered to minimize communication-related CPU stall and overhead bubbles. - ==================================================================== - - There are two additional elements compared to the regular schedule: - - Bootstrap Requests + Release Requests: - - Both can have local failure and need to be consensus on. PP needs to guarantee eventual consistency of local failure and flush malfunc requests out as soft error. - - """ - 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_bootstrapped_rids: Optional[List[str]] = None - transferred_rids: List[str] = [] - release_rids: Optional[List[str]] = None - tmbs = [None] * self.pp_loop_size - bmbs = [None] * self.pp_loop_size - - send_req_work = [] - send_bootstrapped_work = [] - send_consensus_bootstrapped_work = [] - send_proxy_work = [] - send_output_work = [] - send_release_work = [] - send_transfer_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_release_rids = None - next_consensus_bootstrapped_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) - - bootstrapped_rids = self._pp_pd_get_bootstrapped_ids() - bmbs[mb_id] = bootstrapped_rids - self._pp_commit_comm_work(send_bootstrapped_work) - - transferred_rids = self._pp_pd_get_prefill_transferred_ids() - self._pp_commit_comm_work(send_transfer_work) - tmbs[mb_id] = transferred_rids - - self.process_prefill_chunk() - batch = self.get_new_batch_prefill() - if self.require_mlp_sync: - batch = self.prepare_mlp_sync_batch(batch) - 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 = self._pp_recv_proxy_tensors() - - 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) - if self.cur_batch: - result, event = self._pp_launch_batch( - mb_id, pp_proxy_tensors, mb_metadata, last_rank_comm_queue - ) - 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, - ) - ) - send_consensus_bootstrapped_work, consensus_bootstrapped_rids = ( - self._pp_pd_send_consensus_bootstrapped_ids( - bmbs, - next_first_rank_mb_id, - consensus_bootstrapped_rids, - bootstrapped_rids, - ) - ) - send_release_work, release_rids = ( - self._pp_pd_send_consensus_release_ids( - tmbs, next_first_rank_mb_id, release_rids, transferred_rids - ) - ) - - if bmbs[next_mb_id] is not None: - next_consensus_bootstrapped_rids = ( - self._pp_recv_pyobj_from_prev_stage() - ) - next_consensus_bootstrapped_rids = self.process_bootstrapped_queue( - next_consensus_bootstrapped_rids - ) - self._pp_commit_comm_work(send_consensus_bootstrapped_work) - if tmbs[next_mb_id] is not None: - next_release_rids = self._pp_recv_pyobj_from_prev_stage() - self._pp_commit_comm_work(send_release_work) - # post-process the coming microbatch - if mbs[next_mb_id] is not None: - 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 tmbs[next_mb_id] is not None: - self.process_disagg_prefill_inflight_queue(next_release_rids) - if not self.pp_group.is_last_rank: - send_req_work = self._pp_send_pyobj_to_next_stage( - recv_reqs, async_send=True - ) - send_bootstrapped_work = self._pp_send_pyobj_to_next_stage( - bootstrapped_rids, async_send=True - ) - send_transfer_work = self._pp_send_pyobj_to_next_stage( - transferred_rids, async_send=True - ) - if self.cur_batch: - 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_bootstrapped_rids = next_consensus_bootstrapped_rids - - self.running_batch.batch_is_full = False - - # When the server is idle, self-check and re-init some states - if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0: - self.check_memory() - 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: @@ -1179,226 +1188,174 @@ class SchedulerPPMixin: 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 +class ChunkSizePredictor: + """ + Predictor for dynamic chunk size based on quadratic latency model. - # 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 + Models latency as: f(l) = a*l^2 + b*l + c + Predicts next chunk size x such that: f(L+x) - f(L) = target_latency + """ - rmbs = [None] * self.pp_loop_size - pmbs = [None] * self.pp_loop_size - tmbs = [None] * self.pp_loop_size + def __init__(self): + self.quadratic_coeff_a = 0.0 + self.linear_coeff_b = 0.0 + self.constant_coeff_c = 0.0 + self.target_latency: Optional[float] = None + self.is_ready = False - send_req_work = [] + def fit(self, seq_lens: List[int], latencies: List[float]): + """Fit quadratic coefficients f(l) = al^2 + bl + c from data points.""" + L = np.array(seq_lens, dtype=np.float64) + T = np.array(latencies, dtype=np.float64) - # 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 len(L) < 8: + raise ValueError( + f"Not enough data points for quadratic fitting ({len(L)} < 8). " + "Need at least 8 samples with different sequence lengths." ) - 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() + # Build design matrix for f(l) = al^2 + bl + c + X = np.column_stack([L * L, L, np.ones_like(L)]) # [l^2, l, 1] + + try: + coeffs, residuals, rank, s = np.linalg.lstsq(X, T, rcond=None) + if len(coeffs) >= 3: + fitted_a = float(coeffs[0]) # quadratic coefficient + fitted_b = float(coeffs[1]) # linear coefficient + fitted_c = float(coeffs[2]) # constant coefficient + else: + raise ValueError("Failed to fit coefficients: insufficient rank") + except np.linalg.LinAlgError as e: + raise ValueError(f"Failed to fit f(l) = al^2 + bl + c: {e}") + + # Validate coefficients + if fitted_a <= 0: + raise ValueError( + f"Fitted quadratic coefficient a={fitted_a:.2e} is not positive. " + "Attention has O(n^2) complexity, so a must be positive. " + "Check warmup data quality." + ) + + if fitted_b < 0: + logger.warning( + f"Fitted linear coefficient b={fitted_b:.2e} is negative. Setting b=0." + ) + fitted_b = 0.0 + + self.quadratic_coeff_a = fitted_a + self.linear_coeff_b = fitted_b + self.constant_coeff_c = fitted_c + + logger.info( + f"[ChunkSizePredictor] Fitted coefficients: a={fitted_a:.2e}, " + f"b={fitted_b:.2e}, c={fitted_c:.2e}" + ) + + def set_target_latency(self, base_chunk_size: int): + """Set target latency based on base chunk size: target = f(base_chunk_size) - f(0).""" + + def f(l: float) -> float: + """Total latency function: f(l) = al^2 + bl + c (or bl + c for linear)""" + return ( + self.quadratic_coeff_a * l * l + + self.linear_coeff_b * l + + self.constant_coeff_c + ) + + self.target_latency = f(float(base_chunk_size)) - f(0.0) + + if self.target_latency <= 0: + raise ValueError( + f"Calculated target_latency={self.target_latency:.2f}ms is not positive. " + "Check warmup data quality." + ) + + logger.info( + f"[ChunkSizePredictor] Target latency: {self.target_latency:.2f}ms " + f"(base_chunk_size={base_chunk_size})" + ) + + def predict_next_chunk_size( + self, + history_len: int, + base_chunk_size: int, + page_size: int, + context_len: int, + max_chunk_size: Optional[int] = None, + ) -> Optional[int]: + """ + Predict next chunk size x such that f(history_len + x) - f(history_len) = target_latency. + + Args: + history_len: Current sequence length (L) + base_chunk_size: Base chunk size + page_size: Page size for alignment + context_len: Maximum context length + max_chunk_size: Maximum allowed chunk size (optional) + + Returns: + Predicted chunk size, or None if prediction fails + """ + if not self.is_ready or self.target_latency is None: + return None + + # Handle quadratic model: f(l) = al^2 + bl + c + if self.quadratic_coeff_a <= 0: + return None + + # Solve f(L+x) - f(L) = T + # where f(L) = a*L^2 + b*L + c + # This expands to: ax^2 + (2aL+b)x - T = 0 + # A = a, B = 2aL + b, C = -T + A = self.quadratic_coeff_a + B = 2 * self.quadratic_coeff_a * history_len + self.linear_coeff_b + C = -self.target_latency + + discriminant = B * B - 4 * A * C + + if discriminant < 0: + logger.warning( + f"Discriminant is negative ({discriminant:.2e}). " + f"No real solution for chunk size. L={history_len}, T={self.target_latency:.2f}ms." + ) + return None + + sqrt_discriminant = math.sqrt(discriminant) + calculated_chunk_size_float = (-B + sqrt_discriminant) / (2 * A) + + if calculated_chunk_size_float <= 0: + logger.warning( + f"Calculated chunk size is non-positive ({calculated_chunk_size_float:.2f}). " + f"L={history_len}, T={self.target_latency:.2f}ms." + ) + return None + + # Use a smooth coefficient to reduce the abrupt decrease in chunk size + smooth_coeff = envs.SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR.get() + smoothed_chunk_size = base_chunk_size + smooth_coeff * ( + calculated_chunk_size_float - base_chunk_size + ) + calculated_chunk_size = int(smoothed_chunk_size) + + # Align to page_size (round down to nearest multiple) + alignment_size = max(page_size, 1) + dynamic_chunk_size = (calculated_chunk_size // alignment_size) * alignment_size + + # Ensure aligned size is at least alignment_size + if dynamic_chunk_size < alignment_size: + dynamic_chunk_size = alignment_size + + # Apply constraints + max_allowed = context_len - history_len - 100 # Leave 100 tokens margin + if max_chunk_size is not None: + max_allowed = min(max_allowed, max_chunk_size) + dynamic_chunk_size = min(dynamic_chunk_size, max_allowed) + + # Align again after min operation + dynamic_chunk_size = (dynamic_chunk_size // alignment_size) * alignment_size + + if dynamic_chunk_size < alignment_size: + return None + + return dynamic_chunk_size