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] Move block pool operations to a separate class (#13973)
Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Cody Yu <hao.yu.cody@gmail.com>
Chen Zhang committed
Mar 1, 2025 at 04:53 UTC
28943d36cee711e707c2f5a1fb554f77eb5ede47
3 files changed
+360
-277
tests/v1/core/test_prefix_caching.py
+49
-40
@@ -1,12 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""Compare the with and without prefix caching."""
3
+from typing import List
4
+
5
import pytest
6
7
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
8
from vllm.sampling_params import SamplingParams
9
from vllm.utils import cdiv
10
+from vllm.v1.core.block_pool import BlockPool
11
from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
9
-from vllm.v1.core.kv_cache_utils import KVCacheBlock, hash_block_tokens
12
+from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
13
+ hash_block_tokens)
14
15
16
def make_request(request_id,
@@ -62,14 +66,14 @@ def test_prefill():
66
for block_id in (0, 1, 2):
67
block_tokens = tuple(all_token_ids[block_id * 16:(block_id + 1) * 16])
68
block_hash = hash_block_tokens(parent_block_hash, block_tokens)
65
- assert manager.block_pool[block_id].block_hash == block_hash
66
- assert manager.block_pool[block_id].ref_cnt == 1
69
+ assert manager.block_pool.blocks[block_id].block_hash == block_hash
70
+ assert manager.block_pool.blocks[block_id].ref_cnt == 1
71
parent_block_hash = block_hash.hash_value
72
73
# Check partial/preallocated block metadata
74
for block_id in (3, 4):
71
- assert manager.block_pool[block_id].block_hash is None
72
- assert manager.block_pool[block_id].ref_cnt == 1
75
+ assert manager.block_pool.blocks[block_id].block_hash is None
76
+ assert manager.block_pool.blocks[block_id].ref_cnt == 1
77
78
# Cache hit in the common prefix when the original block is still in use.
79
# Incomplete 1 block (5 tokens)
@@ -86,20 +90,21 @@ def test_prefill():
90
assert block.ref_cnt == 2
91
92
# At this point, we should have 3 free blocks left.
89
- assert manager.free_block_queue.num_free_blocks == 3
93
+ assert manager.block_pool.free_block_queue.num_free_blocks == 3
94
95
manager.free(req0)
96
manager.free(req1)
97
98
# All blocks should be available.
95
- assert manager.free_block_queue.num_free_blocks == 10
99
+ assert manager.block_pool.free_block_queue.num_free_blocks == 10
100
# The order should be
101
# [unallocated (7, 8, 9)]
102
# [unique_req0 (4, 3)]
103
# [unique_req1 (6, 5)]
104
# [common (2, 1, 0)]
105
assert [
102
- b.block_id for b in manager.free_block_queue.get_all_free_blocks()
106
+ b.block_id
107
+ for b in manager.block_pool.free_block_queue.get_all_free_blocks()
108
] == [7, 8, 9, 4, 3, 6, 5, 2, 1, 0]
109
110
# Cache hit in the common prefix when the original block is already free.
@@ -116,12 +121,14 @@ def test_prefill():
121
122
# Although we only have 5 free blocks, we have 8 blocks in
123
# the free block queue due to lazy removal.
119
- assert manager.free_block_queue.num_free_blocks == 5
124
+ assert manager.block_pool.free_block_queue.num_free_blocks == 5
125
assert all([
121
- b.ref_cnt == 0 for b in manager.free_block_queue.get_all_free_blocks()
126
+ b.ref_cnt == 0
127
+ for b in manager.block_pool.free_block_queue.get_all_free_blocks()
128
])
123
- assert len([b
124
- for b in manager.free_block_queue.get_all_free_blocks()]) == 5
129
+ assert len([
130
+ b for b in manager.block_pool.free_block_queue.get_all_free_blocks()
131
+ ]) == 5
132
133
manager.free(req2)
134
@@ -133,9 +140,9 @@ def test_prefill():
140
blocks = manager.allocate_slots(req3, 16 * 9, computed_blocks)
141
# This block ID order also checks the eviction order.
142
assert [b.block_id for b in blocks] == [9, 4, 3, 6, 5, 8, 7, 2, 1, 0]
136
- assert manager.free_block_queue.num_free_blocks == 0
137
- assert manager.free_block_queue.free_list_head is None
138
- assert manager.free_block_queue.free_list_tail is None
143
+ assert manager.block_pool.free_block_queue.num_free_blocks == 0
144
+ assert manager.block_pool.free_block_queue.free_list_head is None
145
+ assert manager.block_pool.free_block_queue.free_list_tail is None
146
147
148
def test_decode():
@@ -219,13 +226,14 @@ def test_evict():
226
assert len(blocks) == 3 # 3 full blocks
227
last_token_id += 3 * 16
228
222
- assert manager.free_block_queue.num_free_blocks == 0
229
+ assert manager.block_pool.free_block_queue.num_free_blocks == 0
230
231
manager.free(req0)
232
manager.free(req1)
226
- assert manager.free_block_queue.num_free_blocks == 10
233
+ assert manager.block_pool.free_block_queue.num_free_blocks == 10
234
assert [
228
- b.block_id for b in manager.free_block_queue.get_all_free_blocks()
235
+ b.block_id
236
+ for b in manager.block_pool.free_block_queue.get_all_free_blocks()
237
] == [6, 5, 4, 3, 2, 1, 0, 9, 8, 7]
238
239
# Touch the first 2 blocks.
@@ -235,7 +243,7 @@ def test_evict():
243
assert num_computed_tokens == 2 * 16
244
blocks = manager.allocate_slots(req2, 3, computed_blocks)
245
assert [b.block_id for b in blocks] == [6, 5]
238
- assert manager.free_block_queue.num_free_blocks == 6
246
+ assert manager.block_pool.free_block_queue.num_free_blocks == 6
247
248
249
def test_hash_block_correct_reuse():
@@ -274,7 +282,7 @@ def test_hash_block_correct_reuse():
282
blocks = manager.allocate_slots(req, num_tokens - 1, computed_blocks)
283
assert len(blocks) == 1
284
277
- assert manager.block_pool[blocks[0].block_id].block_hash is None
285
+ assert manager.block_pool.blocks[blocks[0].block_id].block_hash is None
286
287
288
def test_computed_blocks_not_evicted():
@@ -413,13 +421,9 @@ def test_cache_blocks():
421
function of KVCacheManager.
422
"""
423
block_size = 4
416
- manager = KVCacheManager(
417
- block_size=block_size,
424
+ block_pool = BlockPool(
425
num_gpu_blocks=5,
419
- max_model_len=8192,
420
- sliding_window=None,
426
enable_caching=True,
422
- num_preallocate_tokens=0,
427
)
428
# Req:
429
# Block 0: [0, 1, 2, 3]
@@ -430,26 +434,31 @@ def test_cache_blocks():
434
435
# Test that blocks are cached correctly for 2 full blocks from the start.
436
blocks = [KVCacheBlock(block_id=i) for i in range(2)]
437
+ block_hashes: List[BlockHashType] = []
438
434
- manager._cache_full_blocks(
439
+ block_pool.cache_full_blocks(
440
request=req,
436
- blk_start_idx=0,
437
- full_blocks=blocks,
438
- prev_block=None,
441
+ blocks=blocks,
442
+ block_hashes=block_hashes,
443
+ num_cached_blocks=0,
444
+ num_full_blocks=2,
445
+ block_size=block_size,
446
)
447
441
- assert len(manager.cached_block_hash_to_block) == 2
448
+ assert len(block_pool.cached_block_hash_to_block) == 2
449
assert all([block.block_hash is not None for block in blocks])
450
451
# Test that blocks that don't start from the beginning are cached correctly.
445
- blocks = [KVCacheBlock(block_id=2)]
446
- manager._cache_full_blocks(
452
+ blocks += [KVCacheBlock(block_id=2)]
453
+ block_pool.cache_full_blocks(
454
request=req,
448
- blk_start_idx=2,
449
- full_blocks=blocks,
450
- prev_block=None,
455
+ blocks=blocks,
456
+ block_hashes=block_hashes,
457
+ num_cached_blocks=2,
458
+ num_full_blocks=3,
459
+ block_size=block_size,
460
)
452
- assert len(manager.cached_block_hash_to_block) == 3
461
+ assert len(block_pool.cached_block_hash_to_block) == 3
462
assert blocks[0].block_hash is not None
463
464
@@ -580,7 +589,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
589
# Req3 is Req2 + 3 new blocks, so the first 6 blocks are computed,
590
# but it cannot be allocated due to insufficient free blocks (2).
591
# In this case, the ref_cnt of the computed blocks should not be changed.
583
- assert manager.free_block_queue.num_free_blocks == 5
592
+ assert manager.block_pool.free_block_queue.num_free_blocks == 5
593
req3 = make_request("3", common_token_ids * 3)
594
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
595
assert computed_blocks == block_part1
@@ -621,12 +630,12 @@ def test_reset_prefix_cache():
630
631
# Failed to reset prefix cache because some blocks are not freed yet.
632
assert not manager.reset_prefix_cache()
624
- assert manager.cached_block_hash_to_block
633
+ assert manager.block_pool.cached_block_hash_to_block
634
635
# Free the blocks.
636
manager.free(req0)
637
manager.free(req1)
638
639
assert manager.reset_prefix_cache()
631
- assert not manager.cached_block_hash_to_block
632
- assert all([blk.block_hash is None for blk in manager.block_pool])
640
+ assert not manager.block_pool.cached_block_hash_to_block
641
+ assert all([blk.block_hash is None for blk in manager.block_pool.blocks])
vllm/v1/core/block_pool.py
new
+285
@@ -0,0 +1,285 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from collections import defaultdict
3
+from typing import Dict, Iterable, List, Optional
4
+
5
+from vllm.logger import init_logger
6
+from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,
7
+ KVCacheBlock,
8
+ generate_block_hash_extra_keys,
9
+ hash_block_tokens)
10
+from vllm.v1.request import Request
11
+
12
+logger = init_logger(__name__)
13
+
14
+
15
+class BlockPool:
16
+ """BlockPool that manages KVCacheBlocks.
17
+ It provides methods to allocate, free and cache the kv cache blocks. The
18
+ free_block_queue stores the free blocks in eviction order to enable
19
+ allocation, free, and cache eviction. The cached_block_hash_to_block
20
+ maps between block hash and cached block to support finding cached blocks
21
+ by their block hash.
22
+
23
+ Args:
24
+ num_gpu_blocks: The number of blocks in the pool.
25
+ enable_caching: Whether to enable prefix caching.
26
+ """
27
+
28
+ def __init__(self, num_gpu_blocks: int, enable_caching: bool):
29
+ self.num_gpu_blocks = num_gpu_blocks
30
+ self.enable_caching = enable_caching
31
+ # All kv-cache blocks.
32
+ self.blocks: List[KVCacheBlock] = [
33
+ KVCacheBlock(idx) for idx in range(num_gpu_blocks)
34
+ ]
35
+ # Free block queue that constructs and manipulates a doubly linked
36
+ # list of free blocks (including eviction candidates when caching is
37
+ # enabled).
38
+ self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
39
+
40
+ # {block_hash: {block ID: block}}. A cached block is
41
+ # a full block with a block hash that can be used for prefix caching.
42
+ # The cached block may be used by running requests or in the
43
+ # free_block_queue that could potentially be evicted.
44
+ # NOTE: We currently don't de-duplicate the blocks in the cache,
45
+ # meaning that if a block becomes full and is cached, we don't check
46
+ # if there is already an identical block in the cache. This is because
47
+ # we want to make sure the allocated block IDs won't change so that
48
+ # block tables are append-only.
49
+ self.cached_block_hash_to_block: Dict[BlockHashType, Dict[
50
+ int, KVCacheBlock]] = defaultdict(dict)
51
+
52
+ def get_cached_block(self,
53
+ block_hash: BlockHashType) -> Optional[KVCacheBlock]:
54
+ """Get a cached block by the block hash, or None if cache miss.
55
+ If there are duplicated blocks, we return the first block in the cache.
56
+
57
+ Args:
58
+ block_hash: The hash value of the block.
59
+
60
+ Returns:
61
+ The cached block if it exists, or None.
62
+ """
63
+ if block_hash in self.cached_block_hash_to_block:
64
+ first_block_id = list(
65
+ self.cached_block_hash_to_block[block_hash].keys())[0]
66
+ return self.cached_block_hash_to_block[block_hash][first_block_id]
67
+ return None
68
+
69
+ def cache_full_blocks(
70
+ self,
71
+ request: Request,
72
+ blocks: List[KVCacheBlock],
73
+ block_hashes: List[BlockHashType],
74
+ num_cached_blocks: int,
75
+ num_full_blocks: int,
76
+ block_size: int,
77
+ ) -> None:
78
+ """Cache a list of full blocks for prefix caching.
79
+ This function takes a list of blocks that will have their block hash
80
+ metadata to be updated and cached. Given a request, it computes the
81
+ block hashes for the blocks starting from `num_cached_blocks` to
82
+ `num_full_blocks`, updating the metadata for each block
83
+ and caching them in the `cached_block_hash_to_block`.
84
+
85
+ Args:
86
+ request: The request to cache the blocks.
87
+ blocks: All blocks in the request.
88
+ block_hashes: Block hashes of the blocks in the request. Note that
89
+ this list may be shorter than the blocks list. In this case the
90
+ missed block hash will be computed in this function.
91
+ num_cached_blocks: The number of blocks that are already cached.
92
+ num_full_blocks: The number of blocks that are full and should
93
+ be cached after this function.
94
+ block_size: Number of tokens in each block.
95
+ """
96
+ if num_cached_blocks == num_full_blocks:
97
+ return
98
+ new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
99
+ assert len(block_hashes) >= num_cached_blocks
100
+ new_block_hashes = block_hashes[num_cached_blocks:]
101
+
102
+ # Update the new blocks with the block hashes through the chain.
103
+ if num_cached_blocks == 0:
104
+ prev_block_hash_value = None
105
+ else:
106
+ prev_block = blocks[num_cached_blocks - 1]
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
126
+ assert blk.block_hash is None
127
+
128
+ if i + offset < len(new_block_hashes):
129
+ # The block hash may already be computed in
130
+ # "get_computed_blocks" if the tokens are not generated by
131
+ # this request (either the prompt tokens or the previously
132
+ # generated tokens with preemption). In this case we simply
133
+ # reuse the block hash.
134
+ block_hash = new_block_hashes[i + offset]
135
+ else:
136
+ # Otherwise compute the block hash and cache it in the request
137
+ # in case it will be preempted in the future.
138
+ start_token_idx = blk_idx * block_size
139
+ end_token_idx = (blk_idx + 1) * block_size
140
+ block_tokens = request.all_token_ids[
141
+ start_token_idx:end_token_idx]
142
+ assert len(block_tokens) == block_size, (
143
+ f"Expected {block_size} tokens, got "
144
+ f"{len(block_tokens)} at {blk_idx}th block for request "
145
+ f"{request.request_id}({request})")
146
+
147
+ # Generate extra keys for multi-modal inputs. Note that since
148
+ # we reach to this branch only when the block is completed with
149
+ # generated tokens, we only need to consider the last mm input.
150
+ extra_keys, _ = generate_block_hash_extra_keys(
151
+ request, start_token_idx, end_token_idx, -1)
152
+
153
+ # Compute the hash of the current block.
154
+ block_hash = hash_block_tokens(prev_block_hash_value,
155
+ block_tokens, extra_keys)
156
+ block_hashes.append(block_hash)
157
+
158
+ # Update and added the full block to the cache.
159
+ blk.block_hash = block_hash
160
+ self.cached_block_hash_to_block[block_hash][blk.block_id] = blk
161
+ prev_block_hash_value = block_hash.hash_value
162
+
163
+ def get_new_blocks(self, num_blocks: int) -> List[KVCacheBlock]:
164
+ """Get new blocks from the free block pool.
165
+
166
+ Note that we do not check block cache in this function.
167
+
168
+ Args:
169
+ num_blocks: The number of blocks to allocate.
170
+
171
+ Returns:
172
+ A list of new block.
173
+ """
174
+ if num_blocks > self.get_num_free_blocks():
175
+ raise ValueError(
176
+ f"Cannot get {num_blocks} free blocks from the pool")
177
+
178
+ ret: List[KVCacheBlock] = []
179
+ idx = 0
180
+ while idx < num_blocks:
181
+ # First allocate blocks.
182
+ curr_block = self.free_block_queue.popleft()
183
+ assert curr_block.ref_cnt == 0
184
+
185
+ # If the block is cached, evict it.
186
+ if self.enable_caching:
187
+ self._maybe_evict_cached_block(curr_block)
188
+
189
+ curr_block.incr_ref()
190
+ ret.append(curr_block)
191
+ idx += 1
192
+
193
+ return ret
194
+
195
+ def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
196
+ """
197
+ If a block is cached in `cached_block_hash_to_block`, we reset its hash
198
+ metadata and evict it from the cache.
199
+
200
+ Args:
201
+ block: The block to evict.
202
+
203
+ Returns:
204
+ True if the block is evicted, False otherwise.
205
+ """
206
+ block_hash = block.block_hash
207
+ if block_hash and block_hash in self.cached_block_hash_to_block:
208
+ block.reset_hash()
209
+ del self.cached_block_hash_to_block[block_hash][block.block_id]
210
+
211
+ if len(self.cached_block_hash_to_block[block_hash]) == 0:
212
+ del self.cached_block_hash_to_block[block_hash]
213
+
214
+ return True
215
+ return False
216
+
217
+ def touch(self, blocks: List[KVCacheBlock]) -> None:
218
+ """Touch a block increases its reference count by 1, and may remove
219
+ the block from the free queue. This is used when a block is hit by
220
+ another request with the same prefix.
221
+
222
+ Args:
223
+ blocks: A list of blocks to touch.
224
+ """
225
+ for block in blocks:
226
+ # ref_cnt=0 means this block is in the free list (i.e. eviction
227
+ # candidate), so remove it.
228
+ if block.ref_cnt == 0:
229
+ self.free_block_queue.remove(block)
230
+ block.incr_ref()
231
+
232
+ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
233
+ """Free a list of blocks. The blocks should be ordered by their
234
+ eviction priority, where the first block will be evicted first.
235
+
236
+ Args:
237
+ ordered_blocks: A list of blocks to free ordered by their eviction
238
+ priority.
239
+ """
240
+ for block in ordered_blocks:
241
+ block.decr_ref()
242
+ if block.ref_cnt == 0:
243
+ self.free_block_queue.append(block)
244
+
245
+ def reset_prefix_cache(self) -> bool:
246
+ """Reset prefix cache. This function may be used in RLHF
247
+ flows to invalid prefix caching after the weights are updated,
248
+ or used for resetting prefix caching status for benchmarking.
249
+
250
+ Returns:
251
+ bool: True if the prefix cache is successfully reset,
252
+ False otherwise.
253
+ """
254
+ num_used_blocks = (self.num_gpu_blocks - self.get_num_free_blocks())
255
+ if num_used_blocks > 0:
256
+ logger.warning(
257
+ "Failed to reset prefix cache because some "
258
+ "blocks (%d) are not freed yet", num_used_blocks)
259
+ return False
260
+
261
+ # Remove all hashes so that no new blocks will hit.
262
+ self.cached_block_hash_to_block = defaultdict(dict)
263
+
264
+ # Remove all hashes from all blocks.
265
+ for block in self.blocks:
266
+ block.reset_hash()
267
+
268
+ logger.info("Successfully reset prefix cache")
269
+ return True
270
+
271
+ def get_num_free_blocks(self) -> int:
272
+ """Get the number of free blocks in the pool.
273
+
274
+ Returns:
275
+ The number of free blocks.
276
+ """
277
+ return self.free_block_queue.num_free_blocks
278
+
279
+ def get_usage(self) -> float:
280
+ """Get the KV cache usage.
281
+
282
+ Returns:
283
+ The KV cache usage (between 0.0 and 1.0).
284
+ """
285
+ return 1.0 - (self.get_num_free_blocks() / self.num_gpu_blocks)
vllm/v1/core/kv_cache_manager.py
+26
-237
@@ -5,10 +5,8 @@ from typing import DefaultDict, Dict, Iterable, List, Optional, Tuple
5
6
from vllm.logger import init_logger
7
from vllm.utils import cdiv
8
-from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,
9
- KVCacheBlock,
10
- generate_block_hash_extra_keys,
11
- hash_block_tokens,
8
+from vllm.v1.core.block_pool import BlockPool
9
+from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
10
hash_request_tokens)
11
from vllm.v1.metrics.stats import PrefixCacheStats
12
from vllm.v1.request import Request, RequestStatus
@@ -49,26 +47,7 @@ class KVCacheManager:
47
self.num_preallocate_tokens = num_preallocate_tokens
48
self.num_preallocate_blocks = cdiv(num_preallocate_tokens, block_size)
49
52
- # A Block pool of all kv-cache blocks.
53
- self.block_pool: List[KVCacheBlock] = [
54
- KVCacheBlock(idx) for idx in range(num_gpu_blocks)
55
- ]
56
- # Free block queue that constructs and manipulates a doubly linked
57
- # list of free blocks (including eviction candidates when caching is
58
- # enabled).
59
- self.free_block_queue = FreeKVCacheBlockQueue(self.block_pool)
60
-
61
- # {block_hash: {block ID: block}}. A cached block is
62
- # a full block with a block hash that can be used for prefix caching.
63
- # The cached block may be used by running requests or in the
64
- # free_block_queue that could potentially be evicted.
65
- # NOTE: We currently don't de-duplicate the blocks in the cache,
66
- # meaning that if a block becomes full and is cached, we don't check
67
- # if there is already an identical block in the cache. This is because
68
- # we want to make sure the allocated block IDs won't change so that
69
- # block tables are append-only.
70
- self.cached_block_hash_to_block: Dict[BlockHashType, Dict[
71
- int, KVCacheBlock]] = defaultdict(dict)
50
+ self.block_pool = BlockPool(num_gpu_blocks, enable_caching)
51
52
# Mapping from request ID to blocks to track the blocks allocated
53
# for each request, so that we can free the blocks when the request
@@ -96,8 +75,7 @@ class KVCacheManager:
75
Returns:
76
The KV cache usage (between 0.0 and 1.0).
77
"""
99
- return 1.0 - (self.free_block_queue.num_free_blocks /
100
- self.num_gpu_blocks)
78
+ return self.block_pool.get_usage()
79
80
def make_prefix_cache_stats(self) -> PrefixCacheStats:
81
"""Get (and reset) the prefix cache stats.
@@ -139,7 +117,7 @@ class KVCacheManager:
117
# block_hashes is a chain of block hashes. If a block hash is not
118
# in the cached_block_hash_to_id, the following block hashes are
119
# not computed yet for sure.
142
- if cached_block := self._get_cached_block(block_hash):
120
+ if cached_block := self.block_pool.get_cached_block(block_hash):
121
computed_blocks.append(cached_block)
122
else:
123
break
@@ -204,14 +182,14 @@ class KVCacheManager:
182
# when allocating this request.
183
num_evictable_computed_blocks = sum(1 for blk in new_computed_blocks
184
if blk.ref_cnt == 0)
207
- if (num_new_blocks > self.free_block_queue.num_free_blocks -
185
+ if (num_new_blocks > self.block_pool.get_num_free_blocks() -
186
num_evictable_computed_blocks):
187
# Cannot allocate new blocks
188
return None
189
190
# Touch the computed blocks to make sure they won't be evicted.
191
if self.enable_caching:
214
- self._touch(new_computed_blocks)
192
+ self.block_pool.touch(new_computed_blocks)
193
else:
194
assert not new_computed_blocks, (
195
"Computed blocks should be empty when "
@@ -231,7 +209,7 @@ class KVCacheManager:
209
# preallocated blocks.
210
num_new_blocks = min(
211
num_new_blocks + self.num_preallocate_blocks,
234
- self.free_block_queue.num_free_blocks,
212
+ self.block_pool.get_num_free_blocks(),
213
# Should not exceed the maximum number of blocks per request.
214
# This is especially because the block table has the shape
215
# [..., max_num_blocks_per_req].
@@ -240,29 +218,30 @@ class KVCacheManager:
218
assert num_new_blocks > 0
219
220
# Concatenate the computed block IDs and the new block IDs.
243
- new_blocks = self._get_new_blocks(num_new_blocks)
221
+ new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
222
req_blocks.extend(new_blocks)
223
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]
230
# Speculated tokens might be rejected in the future, so we does
231
# not cache any speculated tokens. We only cache blocks with
232
# generated (accepted) tokens.
233
num_full_blocks_after_append = (num_computed_tokens + num_tokens - len(
234
request.spec_token_ids)) // self.block_size
255
- new_full_blocks = req_blocks[
256
- num_cached_blocks:num_full_blocks_after_append]
257
-
258
- if new_full_blocks:
259
- self._cache_full_blocks(
260
- request=request,
261
- blk_start_idx=num_cached_blocks,
262
- # The new full blocks are the full blocks that are not computed.
263
- full_blocks=new_full_blocks,
264
- prev_block=(req_blocks[num_cached_blocks -
265
- 1] if num_cached_blocks > 0 else None))
235
+
236
+ self.block_pool.cache_full_blocks(
237
+ request=request,
238
+ blocks=req_blocks,
239
+ block_hashes=self.req_to_block_hashes[request.request_id],
240
+ num_cached_blocks=num_cached_blocks,
241
+ num_full_blocks=num_full_blocks_after_append,
242
+ block_size=self.block_size,
243
+ )
244
+
245
self.num_cached_block[
246
request.request_id] = num_full_blocks_after_append
247
return new_blocks
@@ -283,11 +262,7 @@ class KVCacheManager:
262
# freed first.
263
ordered_blocks = reversed(blocks)
264
286
- for block in ordered_blocks:
287
- block.decr_ref()
288
- if block.ref_cnt == 0:
289
- self.free_block_queue.append(block)
290
-
265
+ self.block_pool.free_blocks(ordered_blocks)
266
self.num_cached_block.pop(request.request_id, None)
267
268
def reset_prefix_cache(self) -> bool:
@@ -299,25 +274,10 @@ class KVCacheManager:
274
bool: True if the prefix cache is successfully reset,
275
False otherwise.
276
"""
302
- num_used_blocks = (self.num_gpu_blocks -
303
- self.free_block_queue.num_free_blocks)
304
- if num_used_blocks > 0:
305
- logger.warning(
306
- "Failed to reset prefix cache because some "
307
- "blocks (%d) are not freed yet", num_used_blocks)
308
- return False
309
-
310
- # Remove all hashes so that no new blocks will hit.
311
- self.cached_block_hash_to_block = defaultdict(dict)
312
-
313
- # Remove all hashes from all blocks.
314
- for block in self.block_pool:
315
- block.reset_hash()
316
-
317
- self.prefix_cache_stats.reset = True
318
-
319
- logger.info("Successfully reset prefix cache")
320
- return True
277
+ if self.block_pool.reset_prefix_cache():
278
+ self.prefix_cache_stats.reset = True
279
+ return True
280
+ return False
281
282
def get_num_common_prefix_blocks(
283
self,
@@ -367,177 +327,6 @@ class KVCacheManager:
327
break
328
return num_common_blocks
329
370
- def _get_new_blocks(self, num_blocks: int) -> List[KVCacheBlock]:
371
- """Get new blocks from the free block pool.
372
-
373
- Note that we do not check block cache in this function.
374
-
375
- Args:
376
- num_blocks: The number of blocks to allocate.
377
-
378
- Returns:
379
- A list of new block.
380
- """
381
- if num_blocks > self.free_block_queue.num_free_blocks:
382
- raise ValueError(
383
- f"Cannot get {num_blocks} free blocks from the pool")
384
-
385
- ret: List[KVCacheBlock] = []
386
- idx = 0
387
- while idx < num_blocks:
388
- # First allocate blocks.
389
- curr_block = self.free_block_queue.popleft()
390
- assert curr_block.ref_cnt == 0
391
-
392
- # If the block is cached, evict it.
393
- if self.enable_caching:
394
- self._maybe_evict_cached_block(curr_block)
395
-
396
- curr_block.incr_ref()
397
- ret.append(curr_block)
398
- idx += 1
399
-
400
- return ret
401
-
402
- def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
403
- """
404
- If a block is cached in `cached_block_hash_to_block`, we reset its hash
405
- metadata and evict it from the cache.
406
-
407
- Args:
408
- block: The block to evict.
409
-
410
- Returns:
411
- True if the block is evicted, False otherwise.
412
- """
413
- block_hash = block.block_hash
414
- if block_hash and block_hash in self.cached_block_hash_to_block:
415
- block.reset_hash()
416
- del self.cached_block_hash_to_block[block_hash][block.block_id]
417
-
418
- if len(self.cached_block_hash_to_block[block_hash]) == 0:
419
- del self.cached_block_hash_to_block[block_hash]
420
-
421
- return True
422
- return False
423
-
424
- def _get_cached_block(self,
425
- block_hash: BlockHashType) -> Optional[KVCacheBlock]:
426
- """Get a cached block by the block hash, or None if cache miss.
427
- If there are duplicated blocks, we return the first block in the cache.
428
-
429
- Args:
430
- block_hash: The hash value of the block.
431
-
432
- Returns:
433
- The cached block if it exists, or None.
434
- """
435
- if block_hash in self.cached_block_hash_to_block:
436
- first_block_id = list(
437
- self.cached_block_hash_to_block[block_hash].keys())[0]
438
- return self.cached_block_hash_to_block[block_hash][first_block_id]
439
- return None
440
-
441
- def _touch(self, blocks: List[KVCacheBlock]) -> None:
442
- """Touch a block increases its reference count by 1, and may remove
443
- the block from the free queue. This is used when a block is hit by
444
- another request with the same prefix.
445
-
446
- Args:
447
- blocks: A list of blocks to touch.
448
- """
449
- for block in blocks:
450
- # ref_cnt=0 means this block is in the free list (i.e. eviction
451
- # candidate), so remove it.
452
- if block.ref_cnt == 0:
453
- self.free_block_queue.remove(block)
454
- block.incr_ref()
455
-
456
- def _cache_full_blocks(
457
- self,
458
- request: Request,
459
- blk_start_idx: int,
460
- full_blocks: List[KVCacheBlock],
461
- prev_block: Optional[KVCacheBlock],
462
- ) -> None:
463
- """Cache a list of full blocks for prefix caching.
464
-
465
- This function takes a list of blocks that will have their block hash
466
- metadata to be updated and cached. Given a request, it computes the
467
- block hashes for the blocks starting from `blk_start_idx` to the end
468
- of the request's full blocks, updating the metadata for each block
469
- and caching them in the `cached_block_hash_to_block`.
470
-
471
- Args:
472
- request: The request to cache the blocks.
473
- blk_start_idx: The index of the first block in the request's blocks
474
- to cache.
475
- full_blocks: The list of blocks to update hash metadata.
476
- prev_block: The previous block in the chain.
477
- """
478
- block_hashes = self.req_to_block_hashes[request.request_id]
479
- num_cached_block_hashes = len(block_hashes)
480
-
481
- # Update the new blocks with the block hashes through the chain.
482
- prev_block_hash_value = None
483
- if prev_block is not None:
484
- # Previous block must have a block hash because it must be
485
- # a full, cached block.
486
- assert prev_block.block_hash is not None
487
- prev_block_hash_value = prev_block.block_hash.hash_value
488
-
489
- # Find the first uncached block. This case should only happen when
490
- # speculative decoding is used.
491
- offset = 0
492
- for blk in full_blocks:
493
- if blk.block_hash is None:
494
- break
495
- else:
496
- prev_block_hash_value = blk.block_hash.hash_value
497
- offset += 1
498
- else:
499
- # All blocks are cached.
500
- return
501
-
502
- for i, blk in enumerate(full_blocks[offset:]):
503
- blk_idx = blk_start_idx + offset + i
504
- assert blk.block_hash is None
505
-
506
- if blk_idx < num_cached_block_hashes:
507
- # The block hash may already be computed in
508
- # "get_computed_blocks" if the tokens are not generated by
509
- # this request (either the prompt tokens or the previously
510
- # generated tokens with preemption). In this case we simply
511
- # reuse the block hash.
512
- block_hash = block_hashes[blk_idx]
513
- else:
514
- # Otherwise compute the block hash and cache it in the request
515
- # in case it will be preempted in the future.
516
- start_token_idx = blk_idx * self.block_size
517
- end_token_idx = (blk_idx + 1) * self.block_size
518
- block_tokens = request.all_token_ids[
519
- start_token_idx:end_token_idx]
520
- assert len(block_tokens) == self.block_size, (
521
- f"Expected {self.block_size} tokens, got "
522
- f"{len(block_tokens)} at {blk_idx}th block for request "
523
- f"{request.request_id}({request})")
524
-
525
- # Generate extra keys for multi-modal inputs. Note that since
526
- # we reach to this branch only when the block is completed with
527
- # generated tokens, we only need to consider the last mm input.
528
- extra_keys, _ = generate_block_hash_extra_keys(
529
- request, start_token_idx, end_token_idx, -1)
530
-
531
- # Compute the hash of the current block.
532
- block_hash = hash_block_tokens(prev_block_hash_value,
533
- block_tokens, extra_keys)
534
- block_hashes.append(block_hash)
535
-
536
- # Update and added the full block to the cache.
537
- blk.block_hash = block_hash
538
- self.cached_block_hash_to_block[block_hash][blk.block_id] = blk
539
- prev_block_hash_value = block_hash.hash_value
540
-
330
def free_block_hashes(self, request: Request) -> None:
331
"""Discard the block hashes for the request.
332