Use dynamically maintained num_waiting_tokens in get_load() (#13203)

Signed-off-by: Peng Wang <peng_wang@linux.alibaba.com>
Co-authored-by: Peng Wang <peng_wang@linux.alibaba.com>

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 = [
+        preallocated_queue = []
-            entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
+
-        ]
+        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
 

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

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 = [
+        can_run_set = set(can_run_list)
-            x for x in self.waiting_queue if x not in 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
+        # Tokens in waiting queue, bootstrap queue, prealloc/retract queue
-        num_tokens += sum(len(req.origin_input_ids) for req in self.waiting_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: