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 15, 02:15 UTC
[Misc] Simplify the prefix caching logic on draft tokens (#20701)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jul 9, 2025 at 14:48 UTC
7c12a765aa2f2a97ebf6b3bc8361b464461832fc
2 files changed
+10
-11
vllm/v1/core/kv_cache_manager.py
+10
-6
@@ -190,7 +190,6 @@ class KVCacheManager:
190
num_new_tokens: int,
191
num_new_computed_tokens: int = 0,
192
new_computed_blocks: Optional[KVCacheBlocks] = None,
193
- num_draft_tokens: int = 0,
193
num_lookahead_tokens: int = 0,
194
delay_cache_blocks: bool = False,
195
) -> Optional[KVCacheBlocks]:
@@ -286,12 +285,17 @@ class KVCacheManager:
285
if not self.enable_caching or delay_cache_blocks:
286
return KVCacheBlocks(new_blocks)
287
289
- # Speculated tokens might be rejected in the future, so we does
290
- # not cache any speculated tokens. We only cache blocks with
291
- # generated (accepted) tokens.
288
+ # NOTE(woosuk): We want to commit (cache) up to num_computed_tokens +
289
+ # num_new_tokens, but must exclude "non-committable" tokens (e.g.,
290
+ # draft tokens that could be rejected). Therefore, we cap the number
291
+ # at `request.num_tokens`, ensuring only "finalized" tokens are cached.
292
+ num_tokens_to_cache = min(num_computed_tokens + num_new_tokens,
293
+ request.num_tokens)
294
self.coordinator.cache_blocks(
293
- request, self.req_to_block_hashes[request.request_id],
294
- num_computed_tokens + num_new_tokens - num_draft_tokens)
295
+ request,
296
+ self.req_to_block_hashes[request.request_id],
297
+ num_tokens_to_cache,
298
+ )
299
300
return KVCacheBlocks(new_blocks)
301
vllm/v1/core/sched/scheduler.py
-5
@@ -241,15 +241,10 @@ class Scheduler(SchedulerInterface):
241
req_index += 1
242
continue
243
244
- num_draft_tokens = max(
245
- num_new_tokens + request.num_computed_tokens -
246
- request.num_tokens, 0)
247
-
244
while True:
245
new_blocks = self.kv_cache_manager.allocate_slots(
246
request,
247
num_new_tokens,
252
- num_draft_tokens=num_draft_tokens,
248
num_lookahead_tokens=self.num_lookahead_tokens)
249
if new_blocks is None:
250
# The request cannot be scheduled.