Fix scheduler admission for near-full KV requests (#25126)
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user