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, 08:30 UTC
[BugFix][Spec Decode] Improve Prefix Caching Logic in Speculative Decoding (#18668)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
May 24, 2025 at 17:33 UTC
6825d9a998df3f22b5e19fb600d0a0c09950db28
2 files changed
+7
-1
vllm/v1/core/kv_cache_manager.py
+2
-1
@@ -174,6 +174,7 @@ class KVCacheManager:
174
num_new_tokens: int,
175
num_new_computed_tokens: int = 0,
176
new_computed_blocks: Optional[KVCacheBlocks] = None,
177
+ num_draft_tokens: int = 0,
178
num_lookahead_tokens: int = 0,
179
delay_cache_blocks: bool = False,
180
) -> Optional[KVCacheBlocks]:
@@ -273,7 +274,7 @@ class KVCacheManager:
274
# generated (accepted) tokens.
275
self.single_type_manager.cache_blocks(
276
request, self.req_to_block_hashes[request.request_id],
276
- num_computed_tokens + num_new_tokens - len(request.spec_token_ids))
277
+ num_computed_tokens + num_new_tokens - num_draft_tokens)
278
279
return KVCacheBlocks(new_blocks)
280
vllm/v1/core/sched/scheduler.py
+5
@@ -227,10 +227,15 @@ class Scheduler(SchedulerInterface):
227
req_index += 1
228
continue
229
230
+ num_draft_tokens = max(
231
+ num_new_tokens + request.num_computed_tokens -
232
+ request.num_tokens, 0)
233
+
234
while True:
235
new_blocks = self.kv_cache_manager.allocate_slots(
236
request,
237
num_new_tokens,
238
+ num_draft_tokens=num_draft_tokens,
239
num_lookahead_tokens=self.num_lookahead_tokens)
240
if new_blocks is None:
241
# The request cannot be scheduled.