chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Peng Wang <peng_wang@linux.alibaba.com>
Co-authored-by: Peng Wang <peng_wang@linux.alibaba.com>
This commit is contained in:
vipwangerxiao
2025-11-26 20:34:45 +08:00
committed by GitHub
parent 21b0582d4b
commit 15729dbc8e
3 changed files with 40 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/sglang/srt/disaggregation/decode.py
View File

@@ -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

3
python/sglang/srt/disaggregation/prefill.py
View File

@@ -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

35
python/sglang/srt/managers/scheduler.py
View File

@@ -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: