Add page_size to admission token budget check (#22495)

This commit is contained in:
Liangsheng Yin
2026-04-10 01:16:04 -07:00
committed by GitHub
parent dd41764487
commit 6cf7f210bf
+12 -3
View File
@@ -751,9 +751,18 @@ class PrefillAdder:
if req.sampling_params.ignore_eos and getattr(self.tree_cache, "disable", True):
return self.add_one_req_ignore_eos(req)
total_tokens = req.extend_input_len + min(
max(req.sampling_params.max_new_tokens - len(req.output_ids), 0),
CLIP_MAX_NEW_TOKENS,
# Reserve page_size for page-alignment overhead. The paged allocator
# may consume up to one extra page per request (see alloc_extend), and
# _update_prefill_budget already accounts for this in the deduction.
# Without this, admission is more optimistic than the actual budget
# deduction, allowing over-admission when the pool is nearly full.
total_tokens = (
req.extend_input_len
+ min(
max(req.sampling_params.max_new_tokens - len(req.output_ids), 0),
CLIP_MAX_NEW_TOKENS,
)
+ self.page_size
)
# adjusting the input_tokens based on host_hit_length and page_size