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, 21:00 UTC
[v1][Bugfix] Only cache blocks that are not in the prefix cache (#14073)
Chen Zhang committed
Mar 1, 2025 at 16:25 UTC
b9f1d4294e30b700dcb25390c74831a5c178f5fd
2 files changed
+9
-22
vllm/v1/core/block_pool.py
+4
-18
@@ -107,34 +107,20 @@ class BlockPool:
107
assert prev_block.block_hash is not None
108
prev_block_hash_value = prev_block.block_hash.hash_value
109
110
- # Find the first uncached block.
111
- # FIXME: num_cached_blocks should be corrected by the caller
112
- # so this should never happen.
113
- offset = 0
114
- for blk in new_full_blocks:
115
- if blk.block_hash is None:
116
- break
117
- else:
118
- prev_block_hash_value = blk.block_hash.hash_value
119
- offset += 1
120
- else:
121
- # All blocks are cached.
122
- return
123
-
124
- for i, blk in enumerate(new_full_blocks[offset:]):
125
- blk_idx = num_cached_blocks + offset + i
110
+ for i, blk in enumerate(new_full_blocks):
111
assert blk.block_hash is None
112
128
- if i + offset < len(new_block_hashes):
113
+ if i < len(new_block_hashes):
114
# The block hash may already be computed in
115
# "get_computed_blocks" if the tokens are not generated by
116
# this request (either the prompt tokens or the previously
117
# generated tokens with preemption). In this case we simply
118
# reuse the block hash.
134
- block_hash = new_block_hashes[i + offset]
119
+ block_hash = new_block_hashes[i]
120
else:
121
# Otherwise compute the block hash and cache it in the request
122
# in case it will be preempted in the future.
123
+ blk_idx = num_cached_blocks + i
124
start_token_idx = blk_idx * block_size
125
end_token_idx = (blk_idx + 1) * block_size
126
block_tokens = request.all_token_ids[
vllm/v1/core/kv_cache_manager.py
+5
-4
@@ -65,7 +65,7 @@ class KVCacheManager:
65
# This is used to track the number of cached blocks for each request.
66
# This is only used to track the RUNNING requests, we do not track the
67
# data for reempted ones.
68
- self.num_cached_block: Dict[str, int] = defaultdict(int)
68
+ self.num_cached_block: Dict[str, int] = {}
69
self.prefix_cache_stats = PrefixCacheStats()
70
71
@property
@@ -224,9 +224,10 @@ class KVCacheManager:
224
if not self.enable_caching:
225
return new_blocks
226
227
- # FIXME: `num_cached_blocks` is not correct when the prefix cache
228
- # of a new request is hit.
229
- num_cached_blocks = self.num_cached_block[request.request_id]
227
+ # Use `new_computed_blocks` for a new request, and `num_cached_block`
228
+ # for a running request.
229
+ num_cached_blocks = self.num_cached_block.get(request.request_id,
230
+ len(new_computed_blocks))
231
# Speculated tokens might be rejected in the future, so we does
232
# not cache any speculated tokens. We only cache blocks with
233
# generated (accepted) tokens.