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 6, 03:15 UTC
[Core] add option to schedule requests based on full ISL (#37307)
Signed-off-by: Dan Blanaru <48605845+DanBlanaru@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
Dan Blanaru committed
Mar 24, 2026 at 18:01 UTC
dc78c2c933ac8903d7e34bff7b5161eae70595fc
4 files changed
+69
vllm/config/scheduler.py
+6
@@ -135,6 +135,12 @@ class SchedulerConfig:
135
and starting configuration.
136
"""
137
138
+ scheduler_reserve_full_isl: bool = True
139
+ """If True, the scheduler checks whether the full input sequence length
140
+ fits in the KV cache before admitting a new request, rather than only
141
+ checking the first chunk. Prevents over-admission and KV cache thrashing
142
+ with chunked prefill."""
143
+
144
async_scheduling: bool | None = Field(default=None)
145
"""If set to False, disable async scheduling. Async scheduling helps to
146
avoid gaps in GPU utilization, leading to better latency and throughput.
vllm/engine/arg_utils.py
+7
@@ -531,6 +531,8 @@ class EngineArgs:
531
enable_chunked_prefill: bool | None = None
532
disable_chunked_mm_input: bool = SchedulerConfig.disable_chunked_mm_input
533
534
+ scheduler_reserve_full_isl: bool = SchedulerConfig.scheduler_reserve_full_isl
535
+
536
disable_hybrid_kv_cache_manager: bool | None = (
537
SchedulerConfig.disable_hybrid_kv_cache_manager
538
)
@@ -1234,6 +1236,10 @@ class EngineArgs:
1236
scheduler_group.add_argument(
1237
"--scheduler-cls", **scheduler_kwargs["scheduler_cls"]
1238
)
1239
+ scheduler_group.add_argument(
1240
+ "--scheduler-reserve-full-isl",
1241
+ **scheduler_kwargs["scheduler_reserve_full_isl"],
1242
+ )
1243
scheduler_group.add_argument(
1244
"--disable-hybrid-kv-cache-manager",
1245
**scheduler_kwargs["disable_hybrid_kv_cache_manager"],
@@ -1810,6 +1816,7 @@ class EngineArgs:
1816
max_num_partial_prefills=self.max_num_partial_prefills,
1817
max_long_partial_prefills=self.max_long_partial_prefills,
1818
long_prefill_token_threshold=self.long_prefill_token_threshold,
1819
+ scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
1820
disable_hybrid_kv_cache_manager=self.disable_hybrid_kv_cache_manager,
1821
async_scheduling=self.async_scheduling,
1822
stream_interval=self.stream_interval,
vllm/v1/core/kv_cache_manager.py
+39
@@ -215,6 +215,45 @@ class KVCacheManager:
215
216
return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens
217
218
+ def can_fit_full_sequence(
219
+ self,
220
+ request: Request,
221
+ num_new_computed_tokens: int = 0,
222
+ new_computed_blocks: KVCacheBlocks | None = None,
223
+ num_external_computed_tokens: int = 0,
224
+ num_encoder_tokens: int = 0,
225
+ ) -> bool:
226
+ """Check if the KV cache has enough free blocks to hold the full
227
+ sequence, accounting for prefix cache hits and sliding window.
228
+
229
+ This is used as an admission gate to prevent over-admitting requests
230
+ when chunked prefill would otherwise only check the first chunk.
231
+ """
232
+ if new_computed_blocks is not None:
233
+ new_computed_block_list = new_computed_blocks.blocks
234
+ else:
235
+ new_computed_block_list = self.empty_kv_cache_blocks.blocks
236
+
237
+ num_local_computed_tokens = (
238
+ request.num_computed_tokens + num_new_computed_tokens
239
+ )
240
+ total_computed_tokens = min(
241
+ num_local_computed_tokens + num_external_computed_tokens,
242
+ self.max_model_len,
243
+ )
244
+ full_num_tokens = min(request.num_tokens, self.max_model_len)
245
+
246
+ num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
247
+ request_id=request.request_id,
248
+ num_tokens=full_num_tokens,
249
+ new_computed_blocks=new_computed_block_list,
250
+ num_encoder_tokens=num_encoder_tokens,
251
+ total_computed_tokens=total_computed_tokens,
252
+ num_tokens_main_model=full_num_tokens,
253
+ )
254
+
255
+ return num_blocks_to_allocate <= self.block_pool.get_num_free_blocks()
256
+
257
def allocate_slots(
258
self,
259
request: Request,
vllm/v1/core/sched/scheduler.py
+17
@@ -236,6 +236,9 @@ class Scheduler(SchedulerInterface):
236
)
237
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
238
self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
239
+ self.scheduler_reserve_full_isl = (
240
+ self.scheduler_config.scheduler_reserve_full_isl
241
+ )
242
243
self.has_mamba_layers = kv_cache_config.has_mamba_layers
244
self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing
@@ -719,6 +722,20 @@ class Scheduler(SchedulerInterface):
722
for i in encoder_inputs_to_schedule
723
)
724
725
+ if (
726
+ self.scheduler_reserve_full_isl
727
+ and not self.kv_cache_manager.can_fit_full_sequence(
728
+ request,
729
+ num_new_computed_tokens=num_new_local_computed_tokens,
730
+ new_computed_blocks=new_computed_blocks,
731
+ num_external_computed_tokens=num_external_computed_tokens,
732
+ num_encoder_tokens=num_encoder_tokens,
733
+ )
734
+ ):
735
+ if request.has_encoder_inputs:
736
+ self.encoder_cache_manager.free(request)
737
+ break
738
+
739
new_blocks = self.kv_cache_manager.allocate_slots(
740
request,
741
num_new_tokens,