from __future__ import annotations import logging import math import time from array import array from collections import defaultdict, deque from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import numpy as np import torch import torch.distributed from tqdm import tqdm from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.utils import poll_and_all_reduce_attn_cp_tp_group from sglang.srt.distributed.parallel_state import P2PWork from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( get_attention_dp_rank, get_attention_dp_size, is_dp_attention_enabled, set_is_extend_in_batch, ) from sglang.srt.managers.overlap_utils import RelayPayload from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.utils import ( GenerationBatchResult, get_logprob_dict_from_result, get_logprob_from_pp_outputs, ) from sglang.srt.model_executor.forward_batch_info import ( ForwardBatch, ForwardMode, PPProxyTensors, ) from sglang.srt.observability.req_time_stats import set_time_batch from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import DynamicGradMode, broadcast_pyobj, point_to_point_pyobj from sglang.srt.utils.common import get_device_module, is_xpu logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler def _pp_can_skip_output_comm(batch: ScheduleBatch) -> bool: """Check if output send/recv can be skipped for this batch.""" return ( envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM.get() and batch is not None and batch.forward_mode == ForwardMode.EXTEND and len(batch.reqs) == 1 and not batch.contains_last_prefill_chunk and not batch.return_logprob ) @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.ps.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.request_receiver.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"): plan = self.get_next_batch_to_run( running_batch=self.running_batch, last_batch=self.last_batch ) self.running_batch = plan.running_batch self.mbs[mb_id] = plan.batch_to_run self.running_mbs[mb_id] = self.running_batch cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] self.cur_batch_for_debug = cur_batch if 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 cur_batch: result, self.launch_event = self._pp_launch_batch( mb_id, cur_batch, 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 cur_batch: self.device_module.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, msg_type="proxy", ) self.pp_outputs = next_pp_outputs # When the server is idle, self-check and re-init some states if server_is_idle: self.on_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.ps.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.request_receiver.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( last_batch=self.last_batch, running_batch=self.running_batch ) prefill_plan = self.get_new_batch_prefill(self.running_batch) batch = prefill_plan.batch_to_run self.running_batch = prefill_plan.running_batch batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(batch) self.mbs[mb_id] = batch self.running_mbs[mb_id] = self.running_batch cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] self.cur_batch_for_debug = cur_batch if 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 cur_batch: result, self.launch_event = self._pp_launch_batch( mb_id, cur_batch, 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 cur_batch: self.device_module.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, msg_type="proxy", ) 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.on_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.ps.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.request_receiver.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 plan = self.get_next_disagg_decode_batch_to_run( running_batch=self.running_batch ) self.running_batch = plan.running_batch batch = plan.batch_to_run self.mbs[mb_id] = batch self.running_mbs[mb_id] = self.running_batch cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] self.cur_batch_for_debug = cur_batch if cur_batch: server_is_idle = False pp_proxy_tensors = None if not 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 cur_batch: result, self.launch_event = self._pp_launch_batch( mb_id, cur_batch, 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 cur_batch and not cur_batch.forward_mode.is_prebuilt(): self.device_module.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, msg_type="proxy", ) 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.on_idle() def init_pp_loop_state(self: Scheduler): self.pp_loop_size: int = self.ps.pp_size + self.server_args.pp_async_batch_depth # In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation. self.require_attn_tp_allgather = ( not self.server_args.enable_dsa_prefill_context_parallel ) 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.Event, PPProxyTensors]] = deque() self.send_req_work = [] self.send_proxy_work = [] self.send_output_work = [] self.launch_event = None self._pp_tensor_dict_inbox: Dict[str, deque[Dict[str, torch.Tensor]]] = ( defaultdict(deque) ) def profile_and_init_predictor(self: Scheduler): """ Profile prefill latency for dynamic chunk sizing. Only runs on PP0 (first rank), then broadcasts data to all ranks. All ranks fit coefficients using the same data. """ seq_lens: List[int] = [] latencies: List[float] = [] if self.pp_group.is_first_rank: model_runner = self.tp_worker.model_runner model_config = model_runner.model_config input_ids_list: List[array[int]] = [] for i in range(128): chunk_size = int( self.chunked_prefill_size * 1.25 - i * (self.chunked_prefill_size * 1.25 // 128) ) if chunk_size <= 0: break input_ids = array( "q", np.random.randint( 0, 10000, size=chunk_size, dtype=np.int64 ).tobytes(), ) input_ids_list.append(input_ids) sampling_params = SamplingParams( temperature=0, max_new_tokens=1, ) # Create and profile requests for i, input_ids in enumerate( tqdm( input_ids_list, desc="Profiling prefill latency for dynamic chunking", ) ): req = Req( rid=str(i), origin_input_text="", origin_input_ids=input_ids, sampling_params=sampling_params, ) req.full_untruncated_fill_ids = req.origin_input_ids req.logprob_start_len = -1 req.set_extend_range( len(req.prefix_indices), len(req.full_untruncated_fill_ids) ) # Prepare batch batch = ScheduleBatch.init_new( [req], self.req_to_token_pool, self.token_to_kv_pool_allocator, self.tree_cache, self.model_config, False, self.spec_algorithm, ) current_seq_len = req.extend_range.end if is_dp_attention_enabled(): # For profiling, we only have one request on PP0 # Set global_num_tokens to indicate this rank has tokens, others have 0 dp_size = get_attention_dp_size() global_num_tokens = [0] * dp_size dp_rank = get_attention_dp_rank() global_num_tokens[dp_rank] = current_seq_len batch.global_num_tokens = global_num_tokens batch.global_num_tokens_for_logprob = global_num_tokens hs = ( getattr(model_config, "hc_hidden_size", None) or model_config.hidden_size ) proxy_tensors = { "hidden_states": torch.zeros( (current_seq_len, hs), dtype=model_config.dtype, device=self.device, ), "residual": torch.zeros( (current_seq_len, model_config.hidden_size), dtype=model_config.dtype, device=self.device, ), } pp_proxy_topk_size = model_runner.get_pp_proxy_topk_size() if pp_proxy_topk_size is not None: proxy_tensors["topk_indices"] = torch.zeros( (current_seq_len, pp_proxy_topk_size), dtype=torch.int32, device=self.device, ) pp_proxy = PPProxyTensors(proxy_tensors) # Measure latency with device synchronization for accurate timing device_module = get_device_module() # Synchronize before starting timing to ensure clean measurement device_module.synchronize() start = time.perf_counter() batch.prepare_for_extend() # Resolve deferred H2D: prepare_for_extend now leaves input_ids=None if batch.input_ids is None and batch.prefill_input_ids_cpu is not None: batch.input_ids = batch.prefill_input_ids_cpu.to( self.device, non_blocking=True ) batch.prefill_input_ids_cpu = None forward_batch = ForwardBatch.init_new(batch, model_runner) set_is_extend_in_batch(batch.forward_mode.is_extend()) _ = model_runner.forward( forward_batch=forward_batch, pp_proxy_tensors=pp_proxy ) # Synchronize after forward to ensure GPU operations complete device_module.synchronize() latency_seconds = time.perf_counter() - start latency_ms = latency_seconds * 1e3 # Convert to milliseconds seq_lens.append(len(input_ids)) latencies.append(latency_ms) # Release KV cache if req.req_pool_idx is not None: kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : req.extend_range.end ] self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free(req) logger.info( f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: " f"seq_lens={seq_lens}, latencies_ms={latencies}" ) if self.ps.attn_tp_size > 1: data_to_sync_tp = [seq_lens, latencies] data_to_sync_tp = broadcast_pyobj( data_to_sync_tp, self.attn_tp_group.rank, self.attn_tp_cpu_group, src=self.attn_tp_group.ranks[0], ) seq_lens, latencies = data_to_sync_tp if self.ps.attn_cp_size > 1: data_to_sync_tp = [seq_lens, latencies] data_to_sync_tp = broadcast_pyobj( data_to_sync_tp, self.attn_cp_group.rank, self.attn_cp_cpu_group, src=self.attn_cp_group.ranks[0], ) # Broadcast data to all ranks if torch.distributed.is_available() and torch.distributed.is_initialized(): data_to_sync = [seq_lens, latencies] self.pp_group.broadcast_object_list(data_to_sync, src=0) seq_lens, latencies = data_to_sync # Quadratic model: f(l) = al^2 + bl + c self.length_predictor = ChunkSizePredictor() self.length_predictor.fit(seq_lens, latencies) self.length_predictor.set_target_latency(self.chunked_prefill_size) self.length_predictor.is_ready = True logger.info( f"[PP Dynamic Chunk] [PP{self.ps.pp_rank}] Predictor ready (quadratic). " f"Target latency: {self.length_predictor.target_latency:.2f}ms" ) def predict_next_chunk_size(self: Scheduler, history_len: int) -> Optional[int]: """ Predict next chunk size dynamically based on current history length. Args: history_len: Current sequence length Returns: Predicted chunk size, or None to use default chunked_prefill_size """ if ( not self.enable_dynamic_chunking or self.length_predictor is None or not self.length_predictor.is_ready ): return None max_chunk_size = self.max_prefill_tokens predicted_size = self.length_predictor.predict_next_chunk_size( history_len=history_len, base_chunk_size=self.chunked_prefill_size, page_size=self.page_size, context_len=self.model_config.context_len, max_chunk_size=max_chunk_size, ) if predicted_size is not None: logger.debug( f"[PP Dynamic Chunk] [PP{self.ps.pp_rank}] Predicted chunk size: " f"{predicted_size} (history_len={history_len})" ) return predicted_size 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[ Optional[PPProxyTensors], Optional[GenerationBatchResult], Optional[torch.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.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: dp_offset = ( self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size ) p2p_work = point_to_point_pyobj( data, self.ps.pp_rank * self.ps.tp_size + dp_offset, self.world_group.cpu_group, self.ps.pp_rank * self.ps.tp_size + dp_offset, ((self.ps.pp_rank + 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, async_send=async_send, ) return p2p_work def _pp_recv_pyobj_from_prev_stage(self: Scheduler): if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: dp_offset = ( self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size ) data = point_to_point_pyobj( [], self.ps.pp_rank * self.ps.tp_size + dp_offset, self.world_group.cpu_group, ((self.ps.pp_rank - 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, self.ps.pp_rank * self.ps.tp_size + dp_offset, ) else: data = None if self.ps.attn_tp_size > 1: data = broadcast_pyobj( data, self.attn_tp_group.rank, self.attn_tp_cpu_group, src=self.attn_tp_group.ranks[0], ) if self.ps.attn_cp_size > 1: data = broadcast_pyobj( data, self.attn_cp_group.rank, self.attn_cp_cpu_group, src=self.attn_cp_group.ranks[0], ) return data def _pp_prepare_tensor_dict( self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch ) -> Dict[str, torch.Tensor]: tensor_dict = { "next_token_ids": result.next_token_ids, } if batch.return_logprob: logprob_dict = get_logprob_dict_from_result(result) tensor_dict = { **tensor_dict, **logprob_dict, } return tensor_dict def _pp_send_dict_to_next_stage( self: Scheduler, tensor_dict: Dict[str, torch.Tensor], async_send: bool = True, msg_type: str = "default", ): # Warn once if using default untyped messages if msg_type == "default": logger.warning_once( "PP send: using default untyped message. " "Consider adding msg_type='proxy' or 'output' to avoid recv conflicts." ) tensor_dict["__msg_type__"] = msg_type p2p_work = [] p2p_work.extend( self.pp_group.send_tensor_dict( tensor_dict=tensor_dict, all_gather_group=( self.attn_tp_group if self.require_attn_tp_allgather else None ), async_send=async_send, ) ) return p2p_work def _pp_recv_typed_dict( self: Scheduler, expected_kind: str = "default", all_gather_group: Optional = None, ) -> Dict[str, torch.Tensor]: """Receive a typed tensor dict, demultiplexing by msg_type. If a message of the wrong kind is received, it's stashed in the queue and we continue receiving until we get the expected kind. """ if expected_kind in self._pp_tensor_dict_inbox: inbox_queue = self._pp_tensor_dict_inbox[expected_kind] if inbox_queue: return inbox_queue.popleft() while True: tensor_dict = self.pp_group.recv_tensor_dict( all_gather_group=all_gather_group ) received_kind = tensor_dict.get("__msg_type__", "default") if received_kind == expected_kind: if received_kind == "default": logger.warning_once( f"PP recv: got default untyped message. Content keys: {tensor_dict.keys()}" "Consider adding msg_type='proxy' or 'output' to avoid recv conflicts." ) return tensor_dict else: logger.debug( f"PP recv: expected {expected_kind}, got {received_kind}, stashing" ) self._pp_tensor_dict_inbox[received_kind].append(tensor_dict) def _pp_recv_proxy_tensors(self: Scheduler) -> Optional[PPProxyTensors]: pp_proxy_tensors = None if not self.pp_group.is_first_rank: pp_proxy_tensors = PPProxyTensors( self._pp_recv_typed_dict( expected_kind="proxy", all_gather_group=( self.attn_tp_group if self.require_attn_tp_allgather else None ), ) ) return pp_proxy_tensors def _pp_recv_dict_from_prev_stage( self: Scheduler, ) -> Dict[str, torch.Tensor]: return self._pp_recv_typed_dict( expected_kind="output", all_gather_group=( self.attn_tp_group if self.require_attn_tp_allgather else None ), ) def _pp_make_skip_output_result( self: Scheduler, batch: ScheduleBatch, mb_metadata: Optional[PPBatchMetadata], ): bs = len(batch.reqs) placeholder = torch.zeros(bs, dtype=torch.int64, device=self.device) # next_pp_outputs = None so non-last ranks skip forwarding # (pp_outputs is None gate). Placeholder carried in # batch_result.next_token_ids for process_batch_result_prefill. batch.output_ids = placeholder batch_result = GenerationBatchResult( logits_output=None, pp_hidden_states_proxy_tensors=None, next_token_ids=placeholder, can_run_cuda_graph=( mb_metadata.can_run_cuda_graph if mb_metadata else False ), skipped_output_comm=True, ) d2h_event = self.device_module.Event() d2h_event.record(self.device_module.current_stream()) return None, batch_result, d2h_event def _pp_prep_batch_result( self: Scheduler, batch: ScheduleBatch, mb_metadata: PPBatchMetadata, pp_outputs: PPProxyTensors, ): from sglang.srt.managers.scheduler import GenerationBatchResult logits_output = None extend_input_len_per_req = None extend_logprob_start_len_per_req = None if batch.return_logprob: ( logits_output, extend_input_len_per_req, extend_logprob_start_len_per_req, ) = get_logprob_from_pp_outputs(pp_outputs) batch.input_ids = pp_outputs["next_token_ids"].to(torch.int64) # PP rank 0 also relays into output_tokens_buf so the next iter's # resolve_forward_inputs finds these tokens for the decode portion # of mixed-chunk batches (which gather via mix_running_indices). self.future_map.stash( batch.req_pool_indices, RelayPayload(bonus_tokens=batch.input_ids) ) output_result = GenerationBatchResult( logits_output=logits_output, pp_hidden_states_proxy_tensors=None, next_token_ids=pp_outputs["next_token_ids"], extend_input_len_per_req=extend_input_len_per_req, extend_logprob_start_len_per_req=extend_logprob_start_len_per_req, can_run_cuda_graph=mb_metadata.can_run_cuda_graph, ) return output_result def _pp_process_batch_result( self: Scheduler, batch: ScheduleBatch, output_result: GenerationBatchResult ): self.process_batch_result(batch, output_result) def _pp_send_output_to_next_stage( self: Scheduler, next_first_rank_mb_id: int, mbs: List[ScheduleBatch], last_rank_comm_queue: deque, pp_outputs: PPProxyTensors | None, ) -> List[P2PWork]: send_output_work = [] if self.pp_group.is_last_rank: # send ready PP output to rank 0 target = mbs[next_first_rank_mb_id] if target is not None: q_event, pp_outputs_to_send = last_rank_comm_queue.popleft() if ( not target.forward_mode.is_prebuilt() and not _pp_can_skip_output_comm(target) ): self.device_module.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, msg_type="output", ) # 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: with torch.profiler.record_function("send_res_dict_to_next_stage"): send_output_work = self._pp_send_dict_to_next_stage( pp_outputs.tensors, async_send=True, msg_type="output", ) return send_output_work def _pp_send_recv_and_preprocess_output_tensors( self: Scheduler, next_first_rank_mb_id: int, next_mb_id: int, mbs: List[ScheduleBatch], mb_metadata: List[PPBatchMetadata], last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]], pp_outputs: PPProxyTensors | None, ) -> Tuple[ Optional[PPProxyTensors], Optional[GenerationBatchResult], Optional[torch.Event], List[P2PWork], ]: next_pp_outputs = None d2h_event = None batch_result = None send_output_work = [] # On CUDA, isend is async: it enqueues to the stream and returns, # so every rank can send first safely. On some backends isend is # effectively blocking and does not return until the peer posts a # matching recv; if every PP rank sends first, all ranks block # waiting for a receiver and the ring deadlocks. Order send/recv # by pp_rank parity (even: send->recv, odd: recv->send) so each # adjacent pair has one sender and one receiver posted at the # same time. # CUDA: send first # XPU: even ranks send first, odd ranks recv first. send_first = (not is_xpu()) or ((self.ps.pp_rank % 2) == 0) def _do_send(): return self._pp_send_output_to_next_stage( next_first_rank_mb_id, mbs, last_rank_comm_queue, pp_outputs, ) def _do_recv(): nonlocal next_pp_outputs, batch_result, d2h_event target = mbs[next_mb_id] if target is None or target.forward_mode.is_prebuilt(): return if _pp_can_skip_output_comm(target): next_pp_outputs, batch_result, d2h_event = ( self._pp_make_skip_output_result(target, mb_metadata[next_mb_id]) ) return 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.schedule_stream) batch_result = self._pp_prep_batch_result( target, mb_metadata[next_mb_id], next_pp_outputs ) d2h_event = self.device_module.Event() d2h_event.record(self.device_module.current_stream()) if send_first: send_output_work = _do_send() _do_recv() else: _do_recv() send_output_work = _do_send() return next_pp_outputs, batch_result, d2h_event, send_output_work def _pp_launch_batch( self: Scheduler, mb_id: int, cur_batch: ScheduleBatch, pp_proxy_tensors: PPProxyTensors, mb_metadata: List[Optional[PPBatchMetadata]], last_rank_comm_queue: deque, ): with torch.profiler.record_function("run_batch"): with self.forward_stream_ctx: self.forward_stream.wait_stream(self.schedule_stream) set_time_batch( cur_batch.reqs, "set_run_batch_cpu_start_time", trace_only=True, ) result = self.run_batch(cur_batch, pp_proxy_tensors) set_time_batch( cur_batch.reqs, "set_run_batch_cpu_end_time", trace_only=True, attrs={"pp_mb_id": mb_id}, ) mb_metadata[mb_id] = PPBatchMetadata( can_run_cuda_graph=result.can_run_cuda_graph, ) event = self.device_module.Event() event.record(self.device_module.current_stream()) if self.pp_group.is_last_rank: # (last rank) buffer the outputs for async batch depth last_rank_comm_queue.append( ( event, PPProxyTensors( self._pp_prepare_tensor_dict(result, cur_batch) ), ) ) return result, event 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_attn_cp_tp_group( [req.disagg_kv_sender if is_send else req.kv_receiver for req in req_queue], self.attn_cp_cpu_group, self.attn_tp_cpu_group, ) rids: List = [] for poll_statuses in poll_statuses_group: rids.append( [ req.rid if is_send else req.req.rid for req, poll in zip(req_queue, polls) if poll in poll_statuses ] ) return tuple(rids) if len(rids) > 1 else rids[0] 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 ) if self.enable_hisparse: for req in released_reqs: self.hisparse_coordinator.admit_request_direct(req) self.waiting_queue.extend(released_reqs) return [req.rid for req in released_reqs] return None 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.""" # Skip the first data point to reduce fitting bias, as the first run is slower without warmup L = np.array(seq_lens[1:], dtype=np.float64) T = np.array(latencies[1:], 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(length: float) -> float: """Total latency function: f(length) = a*length^2 + b*length + c.""" return ( self.quadratic_coeff_a * length * length + self.linear_coeff_b * length + 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 ) # Make sure the dynamic chunk size is at least 1/4 of the base chunk size calculated_chunk_size = max(int(smoothed_chunk_size), base_chunk_size // 4) # Align to page_size (minimum alignment size is 64) alignment_size = max(page_size, 64) 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