Fix scheduler admission for near-full KV requests (#25126)

This commit is contained in:
ziang663
2026-05-13 15:30:49 -07:00
committed by GitHub
parent f9ff5fc154
commit 8438709e9c
+17 -6
View File
@@ -1865,13 +1865,24 @@ class Scheduler(
self.external_corpus_manager.check_pending_load()
def init_req_max_new_tokens(self, req):
req.sampling_params.max_new_tokens = min(
(
req.sampling_params.max_new_tokens
if req.sampling_params.max_new_tokens is not None
else 1 << 30
input_len = len(req.origin_input_ids)
# Keep this bound consistent with PrefillAdder's admission budget:
# ceil_page(input_len) + max_new_tokens + page_size must be strictly
# smaller than max_total_num_tokens. Otherwise a request can be accepted
# into the waiting queue but can never be scheduled, blocking the queue
# and eventually making health checks fail.
paged_input_len = -(-input_len // self.page_size) * self.page_size
req.sampling_params.max_new_tokens = max(
0,
min(
(
req.sampling_params.max_new_tokens
if req.sampling_params.max_new_tokens is not None
else 1 << 30
),
self.max_req_len - input_len - 1,
self.max_total_num_tokens - paged_input_len - self.page_size - 1,
),
self.max_req_len - len(req.origin_input_ids) - 1,
)
def _process_and_broadcast_mm_inputs(