A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 10:00 UTC
[Core] Simplify handling of `scheduler_reserve_full_isl` option (#41064)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Apr 30, 2026 at 18:10 UTC
dd5506a15759abfc28ed1ca2746221c9ab5e1180
4 files changed
+30
-58
tests/v1/core/test_prefix_caching.py
+4
-2
@@ -2568,7 +2568,9 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
2568
prompt_len = 32 * block_size
2569
req = make_request("long", list(range(prompt_len)), block_size, sha256)
2570
2571
- assert manager.can_fit_full_sequence(req)
2571
+ assert (
2572
+ manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is not None
2573
+ )
2574
2575
2576
def test_can_fit_full_sequence_full_attention_still_gates_oversized():
@@ -2619,4 +2621,4 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():
2621
prompt_len = 16 * block_size
2622
req = make_request("oversized", list(range(prompt_len)), block_size, sha256)
2623
2622
- assert not manager.can_fit_full_sequence(req)
2624
+ assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None
tests/v1/e2e/general/test_mamba_prefix_cache.py
+2
@@ -179,6 +179,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
179
num_external_computed_tokens: int = 0,
180
delay_cache_blocks: bool = False,
181
num_encoder_tokens: int = 0,
182
+ full_sequence_must_fit: bool = False,
183
):
184
ret = original_allocate_slots_fn(
185
self,
@@ -190,6 +191,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
191
num_external_computed_tokens,
192
delay_cache_blocks,
193
num_encoder_tokens,
194
+ full_sequence_must_fit,
195
)
196
if cur_step_action is not None:
197
cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[
vllm/v1/core/kv_cache_manager.py
+23
-42
@@ -222,46 +222,6 @@ class KVCacheManager:
222
223
return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens
224
225
- def can_fit_full_sequence(
226
- self,
227
- request: Request,
228
- num_new_computed_tokens: int = 0,
229
- new_computed_blocks: KVCacheBlocks | None = None,
230
- num_external_computed_tokens: int = 0,
231
- num_encoder_tokens: int = 0,
232
- ) -> bool:
233
- """Check if the KV cache has enough free blocks to hold the full
234
- sequence, accounting for prefix cache hits and sliding window.
235
-
236
- This is used as an admission gate to prevent over-admitting requests
237
- when chunked prefill would otherwise only check the first chunk.
238
- """
239
- if new_computed_blocks is not None:
240
- new_computed_block_list = new_computed_blocks.blocks
241
- else:
242
- new_computed_block_list = self.empty_kv_cache_blocks.blocks
243
-
244
- num_local_computed_tokens = (
245
- request.num_computed_tokens + num_new_computed_tokens
246
- )
247
- total_computed_tokens = min(
248
- num_local_computed_tokens + num_external_computed_tokens,
249
- self.max_model_len,
250
- )
251
- full_num_tokens = min(request.num_tokens, self.max_model_len)
252
-
253
- num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
254
- request_id=request.request_id,
255
- num_tokens=full_num_tokens,
256
- new_computed_blocks=new_computed_block_list,
257
- num_encoder_tokens=num_encoder_tokens,
258
- total_computed_tokens=total_computed_tokens,
259
- num_tokens_main_model=full_num_tokens,
260
- apply_admission_cap=True,
261
- )
262
-
263
- return num_blocks_to_allocate <= self.block_pool.get_num_free_blocks()
264
-
225
def allocate_slots(
226
self,
227
request: Request,
@@ -272,6 +232,7 @@ class KVCacheManager:
232
num_external_computed_tokens: int = 0,
233
delay_cache_blocks: bool = False,
234
num_encoder_tokens: int = 0,
235
+ full_sequence_must_fit: bool = False,
236
) -> KVCacheBlocks | None:
237
"""Add slots for a request with new tokens to append.
238
@@ -293,6 +254,10 @@ class KVCacheManager:
254
num_encoder_tokens: The number of encoder tokens to allocate for
255
cross-attention in encoder-decoder models(e.g., Whisper).
256
For decoder-only models, this should be 0.
257
+ full_sequence_must_fit: Only allocate blocks if the KV cache has enough
258
+ free blocks to hold the full sequence, accounting for prefix cache hits
259
+ and sliding window. Used as an admission gate to prevent over-admitting
260
+ requests when chunked prefill would otherwise only check the first chunk
261
262
Blocks layout:
263
```
@@ -366,10 +331,26 @@ class KVCacheManager:
331
num_local_computed_tokens + num_external_computed_tokens,
332
self.max_model_len,
333
)
334
+
335
+ if full_sequence_must_fit:
336
+ # First check and fail if the full request sequence won't fit.
337
+ full_num_tokens = min(request.num_tokens, self.max_model_len)
338
+
339
+ num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
340
+ request_id=request.request_id,
341
+ num_tokens=full_num_tokens,
342
+ new_computed_blocks=new_computed_block_list,
343
+ num_encoder_tokens=num_encoder_tokens,
344
+ total_computed_tokens=total_computed_tokens,
345
+ num_tokens_main_model=full_num_tokens,
346
+ apply_admission_cap=True,
347
+ )
348
+ if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
349
+ return None
350
+
351
num_tokens_main_model = total_computed_tokens + num_new_tokens
352
num_tokens_need_slot = min(
371
- num_tokens_main_model + num_lookahead_tokens,
372
- self.max_model_len,
353
+ num_tokens_main_model + num_lookahead_tokens, self.max_model_len
354
)
355
356
# Free the blocks that are skipped during the attention computation
vllm/v1/core/sched/scheduler.py
+1
-14
@@ -741,20 +741,6 @@ class Scheduler(SchedulerInterface):
741
for i in encoder_inputs_to_schedule
742
)
743
744
- if (
745
- self.scheduler_reserve_full_isl
746
- and not self.kv_cache_manager.can_fit_full_sequence(
747
- request,
748
- num_new_computed_tokens=num_new_local_computed_tokens,
749
- new_computed_blocks=new_computed_blocks,
750
- num_external_computed_tokens=num_external_computed_tokens,
751
- num_encoder_tokens=num_encoder_tokens,
752
- )
753
- ):
754
- if request.has_encoder_inputs:
755
- self.encoder_cache_manager.free(request)
756
- break
757
-
744
new_blocks = self.kv_cache_manager.allocate_slots(
745
request,
746
num_new_tokens,
@@ -764,6 +750,7 @@ class Scheduler(SchedulerInterface):
750
num_external_computed_tokens=num_external_computed_tokens,
751
delay_cache_blocks=load_kv_async,
752
num_encoder_tokens=num_encoder_tokens,
753
+ full_sequence_must_fit=self.scheduler_reserve_full_isl,
754
)
755
756
if new_blocks is None: