diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 038a6dc21..d3c747dce 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -299,6 +299,11 @@ class DecodePreallocQueue: if is_retracted: self.retracted_queue.append(req) + self.scheduler.num_waiting_tokens += ( + len(req.origin_input_ids) + + len(req.output_ids) + + self.num_reserved_decode_tokens + ) else: if req.bootstrap_host == FAKE_BOOTSTRAP_HOST: kv_receiver_class = get_kv_class( @@ -321,6 +326,7 @@ class DecodePreallocQueue: self.queue.append( DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False) ) + self.scheduler.num_waiting_tokens += len(req.origin_input_ids) def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: if len(req.origin_input_ids) > self.max_total_num_tokens: @@ -361,6 +367,8 @@ class DecodePreallocQueue: req.is_retracted = False self._pre_alloc(req) allocatable_tokens -= required_tokens_for_request + # update counter for resumed request + self.scheduler.num_waiting_tokens -= required_tokens_for_request # load from cpu, release the cpu copy req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator) @@ -536,9 +544,15 @@ class DecodePreallocQueue: RequestStage.DECODE_BOOTSTRAP, decode_req.req.rid, auto_next_anon=True ) - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + preallocated_queue = [] + + for i, entry in enumerate(self.queue): + if i not in indices_to_remove: + preallocated_queue.append(entry) + else: + # update counter for allocated or failed request + self.scheduler.num_waiting_tokens -= len(entry.req.origin_input_ids) + self.queue = preallocated_queue return preallocated_reqs diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 952374ed5..da099a465 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -200,6 +200,7 @@ class PrefillBootstrapQueue: self._process_req(req) req.add_latency(RequestStage.PREFILL_PREPARE) self.queue.append(req) + self.scheduler.num_waiting_tokens += len(req.origin_input_ids) trace_slice_end(RequestStage.PREFILL_PREPARE, req.rid, auto_next_anon=True) def extend(self, reqs: List[Req], num_kv_heads: int) -> None: @@ -270,6 +271,8 @@ class PrefillBootstrapQueue: self.scheduler.stream_output([req], req.return_logprob) indices_to_remove.add(i) failed_reqs.append(req) + # update counter for failed bootstrap request + self.scheduler.num_waiting_tokens -= len(req.origin_input_ids) if self.scheduler.enable_metrics: self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() continue diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b0c528b73..66870540b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -427,6 +427,7 @@ class Scheduler( # Init running status self.waiting_queue: List[Req] = [] + self.num_waiting_tokens = 0 # The running decoding batch for continuous batching self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False) # The current forward batch @@ -1464,6 +1465,7 @@ class Scheduler( return self._prefetch_kvcache(req) self.waiting_queue.append(req) + self.num_waiting_tokens += len(req.origin_input_ids) req.time_stats.wait_queue_entry_time = time.perf_counter() trace_slice_end(RequestStage.REQUEST_PROCESS, req.rid, auto_next_anon=True) elif self.disaggregation_mode == DisaggregationMode.PREFILL: @@ -1531,6 +1533,8 @@ class Scheduler( if abort_existing_req: self.waiting_queue.pop(idx) req_to_abort = candidate_req + # update counter for preempted existing request + self.num_waiting_tokens -= len(req_to_abort.origin_input_ids) message = "The request is aborted by a higher priority request." self.send_to_tokenizer.send_output( @@ -1808,9 +1812,15 @@ class Scheduler( for req in can_run_list: req.add_latency(RequestStage.PREFILL_WAITING) - self.waiting_queue = [ - x for x in self.waiting_queue if x not in set(can_run_list) - ] + can_run_set = set(can_run_list) + waiting_queue = [] + for x in self.waiting_queue: + if x in can_run_set: + self.num_waiting_tokens -= len(x.origin_input_ids) + else: + waiting_queue.append(x) + + self.waiting_queue = waiting_queue if adder.preempt_list: for req in adder.preempt_list: self._add_request_to_queue(req) @@ -2233,8 +2243,6 @@ class Scheduler( return if_success def get_load(self, recv_req: GetLoadReqInput = None) -> GetLoadReqOutput: - # TODO(lsyin): use dynamically maintained num_waiting_tokens - if self.is_hybrid: num_tokens_full = ( self.full_tokens_per_layer @@ -2260,21 +2268,15 @@ class Scheduler( - self.tree_cache.evictable_size() ) - # Tokens in waiting queue, bootstrap queue, prealloc queue - num_tokens += sum(len(req.origin_input_ids) for req in self.waiting_queue) + # Tokens in waiting queue, bootstrap queue, prealloc/retract queue + num_tokens += self.num_waiting_tokens num_waiting_reqs = len(self.waiting_queue) if self.disaggregation_mode == DisaggregationMode.PREFILL: - num_tokens += sum( - len(req.origin_input_ids) - for req in self.disagg_prefill_bootstrap_queue.queue - ) num_waiting_reqs += len(self.disagg_prefill_bootstrap_queue.queue) elif self.disaggregation_mode == DisaggregationMode.DECODE: - num_tokens += sum( - len(req.req.origin_input_ids) - for req in self.disagg_decode_prealloc_queue.queue - ) num_waiting_reqs += len(self.disagg_decode_prealloc_queue.queue) + num_waiting_reqs += len(self.disagg_decode_transfer_queue.queue) + num_waiting_reqs += len(self.disagg_decode_prealloc_queue.retracted_queue) return GetLoadReqOutput( dp_rank=self.dp_rank, @@ -2390,6 +2392,9 @@ class Scheduler( # For disaggregation decode mode, the request in the waiting queue has KV cache allocated. if self.disaggregation_mode == DisaggregationMode.DECODE: release_kv_cache(req, self.tree_cache) + else: + # update counter for abort request in prealloc and waiting queue + self.num_waiting_tokens -= len(req.origin_input_ids) # For mamba radix cache if req.mamba_pool_idx is not None: