20231088/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[V1][Minor] Minor simplification for get_computed_blocks (#16139)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2025-04-06 20:38:12 -07:00 committed by GitHub
parent 86fc2321ff
commit 3749e28774
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 32 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
vllm/v1/core/kv_cache_manager.py
View File

@ -126,7 +126,10 @@ class KVCacheManager:
self.req_to_block_hashes[request.request_id] = block_hashes self.req_to_block_hashes[request.request_id] = block_hashes
self.prefix_cache_stats.requests += 1 self.prefix_cache_stats.requests += 1
if request.sampling_params.prompt_logprobs is None: # When the request requires prompt logprobs, we skip prefix caching.
if request.sampling_params.prompt_logprobs is not None:
return [], 0
if len(block_hashes) * self.block_size == request.num_tokens: if len(block_hashes) * self.block_size == request.num_tokens:
# When prompt length is divisible by the block size and all # When prompt length is divisible by the block size and all
# blocks are cached, we need to recompute the last token. This # blocks are cached, we need to recompute the last token. This
@ -142,22 +145,20 @@ class KVCacheManager:
computed_blocks = ( computed_blocks = (
self.specialized_manager.find_longest_cache_hit(block_hashes)) self.specialized_manager.find_longest_cache_hit(block_hashes))
self.prefix_cache_stats.queries += len(block_hashes)
self.prefix_cache_stats.hits += len(computed_blocks)
if last_block_hash is not None: if last_block_hash is not None:
# Add back the last block hash if it was removed. # Add back the last block hash if it was removed.
# NOTE: Because block_hashes is cached in req_to_block_hashes,
# we shouldn't modify it directly.
block_hashes.append(last_block_hash) block_hashes.append(last_block_hash)
self.prefix_cache_stats.queries += len(block_hashes)
self.prefix_cache_stats.hits += len(computed_blocks)
# NOTE(woosuk): Since incomplete blocks are not eligible for # NOTE(woosuk): Since incomplete blocks are not eligible for
# sharing, `num_computed_tokens` is always a multiple of # sharing, `num_computed_tokens` is always a multiple of
# `block_size`. # `block_size`.
num_computed_tokens = len(computed_blocks) * self.block_size num_computed_tokens = len(computed_blocks) * self.block_size
return computed_blocks, num_computed_tokens return computed_blocks, num_computed_tokens
else:
# Skip cache hits for prompt logprobs
return [], 0
def allocate_slots( def allocate_slots(
self, self,