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, 07:15 UTC
[v1][KVCacheManager] Rename BlockHashType to BlockHash (#19015)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Jun 3, 2025 at 16:01 UTC
f32fcd944430603ebcbbf04454b4e15754168ef4
8 files changed
+28
-30
docs/design/v1/prefix_caching.md
+1
-1
@@ -104,7 +104,7 @@ class KVCacheBlock:
104
block_id: int
105
# The block hash (will be assigned when the block is full,
106
# and will be reset when the block is evicted).
107
- block_hash: BlockHashType
107
+ block_hash: BlockHash
108
# The number of requests using this block now.
109
ref_cnt: int
110
tests/v1/core/test_kv_cache_utils.py
+5
-7
@@ -100,8 +100,8 @@ def test_kv_cache_block():
100
assert block.ref_cnt == 0
101
102
# Test block hash setting and resetting
103
- block_hash = vllm.v1.core.kv_cache_utils.BlockHashType(hash_value=123,
104
- token_ids=(1, 2, 3))
103
+ block_hash = vllm.v1.core.kv_cache_utils.BlockHash(hash_value=123,
104
+ token_ids=(1, 2, 3))
105
block.block_hash = block_hash
106
assert block.block_hash == block_hash
107
@@ -282,7 +282,7 @@ def test_hash_block_tokens(hash_fn):
282
283
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
284
curr_block_token_ids, extra_keys)
285
- assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHashType)
285
+ assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHash)
286
assert block_hash.hash_value == hash_fn(
287
(parent_block_hash, curr_block_token_ids, extra_keys))
288
assert block_hash.token_ids == curr_block_token_ids
@@ -306,10 +306,8 @@ def test_hash_request_tokens(hash_fn):
306
block_hashes = hash_request_tokens(hash_fn, block_size, request)
307
308
assert len(block_hashes) == 2
309
- assert isinstance(block_hashes[0],
310
- vllm.v1.core.kv_cache_utils.BlockHashType)
311
- assert isinstance(block_hashes[1],
312
- vllm.v1.core.kv_cache_utils.BlockHashType)
309
+ assert isinstance(block_hashes[0], vllm.v1.core.kv_cache_utils.BlockHash)
310
+ assert isinstance(block_hashes[1], vllm.v1.core.kv_cache_utils.BlockHash)
311
312
# Check the first block
313
assert block_hashes[0].token_ids == (0, 1, 2)
tests/v1/core/test_prefix_caching.py
+2
-2
@@ -12,7 +12,7 @@ from vllm.sampling_params import SamplingParams
12
from vllm.utils import sha256
13
from vllm.v1.core.block_pool import BlockPool
14
from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
15
-from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
15
+from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
16
hash_block_tokens)
17
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
18
KVCacheGroupSpec, SlidingWindowSpec)
@@ -547,7 +547,7 @@ def test_cache_blocks(hash_fn):
547
548
# Test that blocks are cached correctly for 2 full blocks from the start.
549
blocks = [KVCacheBlock(block_id=i) for i in range(2)]
550
- block_hashes: list[BlockHashType] = []
550
+ block_hashes: list[BlockHash] = []
551
552
block_pool.cache_full_blocks(
553
request=req,
tests/v1/core/test_specialized_manager.py
+2
-2
@@ -3,7 +3,7 @@
3
import torch
4
5
from vllm.v1.core.block_pool import BlockPool
6
-from vllm.v1.core.kv_cache_utils import BlockHashType, KVCacheBlock
6
+from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
7
from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
8
from vllm.v1.kv_cache_interface import SlidingWindowSpec
9
@@ -32,7 +32,7 @@ def test_sliding_window_possible_cached_prefix():
32
33
def run_one_case(block_is_cached, expect_length):
34
block_hash_list = [
35
- BlockHashType(i, ()) for i in range(len(block_is_cached))
35
+ BlockHash(i, ()) for i in range(len(block_is_cached))
36
]
37
38
block_pool.cached_block_hash_to_block.clear()
vllm/v1/core/block_pool.py
+4
-4
@@ -6,7 +6,7 @@ from typing import Callable, Optional
6
from vllm.distributed.kv_events import (AllBlocksCleared, BlockRemoved,
7
BlockStored, KVCacheEvent)
8
from vllm.logger import init_logger
9
-from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,
9
+from vllm.v1.core.kv_cache_utils import (BlockHash, FreeKVCacheBlockQueue,
10
KVCacheBlock,
11
generate_block_hash_extra_keys,
12
hash_block_tokens)
@@ -55,7 +55,7 @@ class BlockPool:
55
# if there is already an identical block in the cache. This is because
56
# we want to make sure the allocated block IDs won't change so that
57
# block tables are append-only.
58
- self.cached_block_hash_to_block: dict[BlockHashType, dict[
58
+ self.cached_block_hash_to_block: dict[BlockHash, dict[
59
int, KVCacheBlock]] = defaultdict(dict)
60
61
# To represent a placeholder block with block_id=0.
@@ -67,7 +67,7 @@ class BlockPool:
67
self.kv_event_queue: list[KVCacheEvent] = []
68
69
def get_cached_block(self,
70
- block_hash: BlockHashType) -> Optional[KVCacheBlock]:
70
+ block_hash: BlockHash) -> Optional[KVCacheBlock]:
71
"""Get a cached block by the block hash, or None if cache miss.
72
If there are duplicated blocks, we return the first block in the cache.
73
@@ -87,7 +87,7 @@ class BlockPool:
87
self,
88
request: Request,
89
blocks: list[KVCacheBlock],
90
- block_hashes: list[BlockHashType],
90
+ block_hashes: list[BlockHash],
91
num_cached_blocks: int,
92
num_full_blocks: int,
93
block_size: int,
vllm/v1/core/kv_cache_manager.py
+2
-2
@@ -8,7 +8,7 @@ from vllm.distributed.kv_events import KVCacheEvent
8
from vllm.logger import init_logger
9
from vllm.utils import sha256
10
from vllm.v1.core.block_pool import BlockPool
11
-from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
11
+from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
12
hash_request_tokens)
13
from vllm.v1.core.single_type_kv_cache_manager import (
14
get_manager_for_kv_cache_spec)
@@ -92,7 +92,7 @@ class KVCacheManager:
92
# This is to avoid recomputing the block hashes for each call of
93
# `get_computed_blocks` or `allocate_slots`.
94
self.req_to_block_hashes: defaultdict[
95
- str, list[BlockHashType]] = defaultdict(list)
95
+ str, list[BlockHash]] = defaultdict(list)
96
97
@property
98
def usage(self) -> float:
vllm/v1/core/kv_cache_utils.py
+7
-7
@@ -18,7 +18,7 @@ from vllm.v1.request import Request
18
logger = init_logger(__name__)
19
20
21
-class BlockHashType(NamedTuple):
21
+class BlockHash(NamedTuple):
22
"""Hash value of a block (int), the token IDs in the block, and extra keys.
23
We keep a tuple of token IDs and extra keys to reduce the likelihood of
24
hash collisions when the hash value is the same. By using SHA256 however,
@@ -117,7 +117,7 @@ class KVCacheBlock:
117
ref_cnt: int = 0
118
# The hash of the block composed of (block hash, tuple of token IDs).
119
# It is only available when the block is full.
120
- _block_hash: Optional[BlockHashType] = None
120
+ _block_hash: Optional[BlockHash] = None
121
122
# Used to construct a doubly linked list for free blocks.
123
# These two attributes should only be manipulated by FreeKVCacheBlockQueue.
@@ -131,11 +131,11 @@ class KVCacheBlock:
131
self.ref_cnt -= 1
132
133
@property
134
- def block_hash(self) -> Optional[BlockHashType]:
134
+ def block_hash(self) -> Optional[BlockHash]:
135
return self._block_hash
136
137
@block_hash.setter
138
- def block_hash(self, block_hash: BlockHashType):
138
+ def block_hash(self, block_hash: BlockHash):
139
assert self.block_hash is None, (
140
"The block already has a hash. This should not happen.")
141
self._block_hash = block_hash
@@ -398,7 +398,7 @@ def hash_block_tokens(
398
hash_function: Callable,
399
parent_block_hash: Optional[int],
400
curr_block_token_ids: Sequence[int],
401
- extra_keys: Optional[tuple[Any, ...]] = None) -> BlockHashType:
401
+ extra_keys: Optional[tuple[Any, ...]] = None) -> BlockHash:
402
"""Computes a hash value corresponding to the contents of a block and
403
the contents of the preceding block(s). The hash value is used for
404
prefix caching. We use LRU cache for this function to avoid recomputing
@@ -419,14 +419,14 @@ def hash_block_tokens(
419
parent_block_hash = NONE_HASH
420
421
curr_block_token_ids_tuple = tuple(curr_block_token_ids)
422
- return BlockHashType(
422
+ return BlockHash(
423
hash_function(
424
(parent_block_hash, curr_block_token_ids_tuple, extra_keys)),
425
curr_block_token_ids_tuple, extra_keys)
426
427
428
def hash_request_tokens(hash_function: Any, block_size: int,
429
- request: Request) -> list[BlockHashType]:
429
+ request: Request) -> list[BlockHash]:
430
"""Computes hash values of a chain of blocks given a sequence of
431
token IDs. The hash value is used for prefix caching.
432
vllm/v1/core/single_type_kv_cache_manager.py
+5
-5
@@ -5,7 +5,7 @@ from typing import Callable
5
6
from vllm.utils import cdiv
7
from vllm.v1.core.block_pool import BlockPool
8
-from vllm.v1.core.kv_cache_utils import BlockHashType, KVCacheBlock
8
+from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
9
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheSpec,
10
SlidingWindowSpec)
11
from vllm.v1.request import Request
@@ -133,7 +133,7 @@ class SingleTypeKVCacheManager(ABC):
133
req_blocks.extend(new_blocks)
134
return new_blocks
135
136
- def cache_blocks(self, request: Request, block_hashes: list[BlockHashType],
136
+ def cache_blocks(self, request: Request, block_hashes: list[BlockHash],
137
num_tokens: int) -> None:
138
"""
139
Cache the blocks for the request.
@@ -187,7 +187,7 @@ class SingleTypeKVCacheManager(ABC):
187
raise NotImplementedError
188
189
@abstractmethod
190
- def find_longest_cache_hit(self, block_hashes: list[BlockHashType],
190
+ def find_longest_cache_hit(self, block_hashes: list[BlockHash],
191
max_length: int) -> list[KVCacheBlock]:
192
"""
193
Get the longest cache hit prefix of the blocks that is not longer than
@@ -228,7 +228,7 @@ class SingleTypeKVCacheManager(ABC):
228
229
class FullAttentionManager(SingleTypeKVCacheManager):
230
231
- def find_longest_cache_hit(self, block_hashes: list[BlockHashType],
231
+ def find_longest_cache_hit(self, block_hashes: list[BlockHash],
232
max_length: int) -> list[KVCacheBlock]:
233
computed_blocks: list[KVCacheBlock] = []
234
max_num_blocks = max_length // self.block_size
@@ -280,7 +280,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
280
self.sliding_window_contiguous_blocks += 1
281
self._null_block = block_pool.null_block
282
283
- def find_longest_cache_hit(self, block_hashes: list[BlockHashType],
283
+ def find_longest_cache_hit(self, block_hashes: list[BlockHash],
284
max_length: int) -> list[KVCacheBlock]:
285
# TODO: reduce i by sliding_window_contiguous_blocks when cache miss, to
286
# optimize the time complexity from O(max_num_blocks) to