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, 11:00 UTC
[v1] Move block management logic from KVCacheManager to SpecializedManager (#17474)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
May 9, 2025 at 23:25 UTC
200da9a51751973740f4dc71f0d1e13cc5698cb0
6 files changed
+269
-155
tests/v1/core/test_kv_cache_utils.py
+3
-3
@@ -539,7 +539,7 @@ def test_allocate_with_lookahead():
539
max_model_len=100)
540
blocks = kv_cache_manager.allocate_slots(
541
request,
542
- num_tokens=3,
542
+ num_new_tokens=3,
543
num_lookahead_tokens=2, # Total required: 3+2=5 tokens
544
)
545
assert len(blocks.blocks) == 2 # ceil(5/4)=2 blocks
@@ -550,7 +550,7 @@ def test_allocate_with_lookahead():
550
# required_blocks = ceil((3 + 2) /4) = 2
551
blocks = kv_cache_manager.allocate_slots(
552
request,
553
- num_tokens=3,
553
+ num_new_tokens=3,
554
num_lookahead_tokens=2,
555
)
556
assert len(blocks.blocks) == 2
@@ -561,7 +561,7 @@ def test_allocate_with_lookahead():
561
max_model_len=100)
562
blocks = kv_cache_manager.allocate_slots(
563
request,
564
- num_tokens=3,
564
+ num_new_tokens=3,
565
num_lookahead_tokens=4,
566
)
567
assert len(blocks.blocks) == 2
tests/v1/core/test_prefix_caching.py
+8
-5
@@ -299,7 +299,8 @@ def test_decode():
299
req0.append_output_token_ids(8)
300
new_blocks = manager.allocate_slots(req0, 4)
301
assert new_blocks is not None and len(new_blocks.blocks) == 0
302
- assert manager.req_to_blocks[req0.request_id][-1].block_hash is None
302
+ assert manager.single_type_manager.req_to_blocks[
303
+ req0.request_id][-1].block_hash is None
304
305
# Append slots with allocating a new block.
306
req0.num_computed_tokens = 59
@@ -309,8 +310,10 @@ def test_decode():
310
req0.append_output_token_ids(7)
311
new_blocks = manager.allocate_slots(req0, 19)
312
assert new_blocks is not None and len(new_blocks.blocks) == 1
312
- assert manager.req_to_blocks[req0.request_id][-2].block_hash is not None
313
- assert manager.req_to_blocks[req0.request_id][-1].block_hash is None
313
+ assert manager.single_type_manager.req_to_blocks[
314
+ req0.request_id][-2].block_hash is not None
315
+ assert manager.single_type_manager.req_to_blocks[
316
+ req0.request_id][-1].block_hash is None
317
318
319
def test_evict():
@@ -689,7 +692,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
692
assert not computed_blocks.blocks
693
assert num_computed_tokens == 0
694
manager.allocate_slots(req0, 48, computed_blocks)
692
- block_part0 = manager.req_to_blocks[req0.request_id]
695
+ block_part0 = manager.single_type_manager.req_to_blocks[req0.request_id]
696
697
# | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |
698
req1 = make_request("1", common_token_ids * 2)
@@ -697,7 +700,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
700
assert computed_blocks.blocks == block_part0
701
assert num_computed_tokens == 3 * 16
702
manager.allocate_slots(req1, 48, computed_blocks)
700
- block_part1 = manager.req_to_blocks[req1.request_id]
703
+ block_part1 = manager.single_type_manager.req_to_blocks[req1.request_id]
704
# | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
705
# | Req1-5(F)| ... |
706
manager.free(req1)
tests/v1/core/test_scheduler.py
+8
-5
@@ -812,10 +812,11 @@ def _assert_right_kv_cache_manager(
812
# Make sure the request stats are right.
813
EXPECTED_TOTAL_BLOCKS = num_tokens // block_size
814
for req_id in req_ids:
815
- blocks = scheduler.kv_cache_manager.req_to_blocks[req_id]
815
+ blocks = (scheduler.kv_cache_manager.single_type_manager.
816
+ req_to_blocks[req_id])
817
hashes = scheduler.kv_cache_manager.req_to_block_hashes[req_id]
817
- assert (scheduler.kv_cache_manager.num_cached_block[req_id] ==
818
- EXPECTED_TOTAL_BLOCKS)
818
+ assert (scheduler.kv_cache_manager.single_type_manager.
819
+ num_cached_block[req_id] == EXPECTED_TOTAL_BLOCKS)
820
assert len(blocks) == EXPECTED_TOTAL_BLOCKS
821
assert len(hashes) == EXPECTED_TOTAL_BLOCKS
822
@@ -1195,9 +1196,11 @@ def assert_scheduler_empty(scheduler: Scheduler):
1196
assert len(scheduler.encoder_cache_manager.cached) == 0
1197
1198
# KVCache Manager.
1198
- assert len(scheduler.kv_cache_manager.req_to_blocks) == 0
1199
+ assert len(
1200
+ scheduler.kv_cache_manager.single_type_manager.req_to_blocks) == 0
1201
assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
1200
- assert len(scheduler.kv_cache_manager.num_cached_block) == 0
1202
+ assert len(
1203
+ scheduler.kv_cache_manager.single_type_manager.num_cached_block) == 0
1204
num_free_blocks = (
1205
scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
1206
assert num_free_blocks == (
tests/v1/core/test_specialized_manager.py
+18
-18
@@ -8,6 +8,14 @@ from vllm.v1.core.specialized_manager import SlidingWindowManager
8
from vllm.v1.kv_cache_interface import SlidingWindowSpec
9
10
11
+def get_sliding_window_manager(sliding_window_spec, block_pool):
12
+ return SlidingWindowManager(sliding_window_spec,
13
+ block_pool,
14
+ use_eagle=False,
15
+ num_kv_cache_groups=1,
16
+ caching_hash_fn=lambda x: x)
17
+
18
+
19
def test_sliding_window_possible_cached_prefix():
20
sliding_window_spec = SlidingWindowSpec(
21
block_size=2,
@@ -19,9 +27,7 @@ def test_sliding_window_possible_cached_prefix():
27
)
28
29
block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
22
- manager = SlidingWindowManager(sliding_window_spec,
23
- block_pool,
24
- use_eagle=False)
30
+ manager = get_sliding_window_manager(sliding_window_spec, block_pool)
31
32
def run_one_case(block_is_cached, expect_length):
33
block_hash_list = [
@@ -81,9 +87,7 @@ def test_sliding_window_remove_skipped_blocks():
87
88
block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
89
84
- manager = SlidingWindowManager(sliding_window_spec,
85
- block_pool,
86
- use_eagle=False)
90
+ manager = get_sliding_window_manager(sliding_window_spec, block_pool)
91
92
null_block_id = block_pool.null_block.block_id
93
@@ -104,39 +108,35 @@ def test_sliding_window_remove_skipped_blocks():
108
1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010
109
]
110
block_table = id_to_block_table(original_block_ids)
107
- removed = manager.remove_skipped_blocks(block_table, 0)
108
- assert_block_id(removed, [])
111
+ manager.req_to_blocks["test"] = block_table
112
+
113
+ manager.remove_skipped_blocks("test", 0)
114
assert_block_id(block_table, original_block_ids)
115
116
# 4 tokens are computed. Only token 0 is out of the sliding window. As
117
# block 1000 also contains token 1 that is in the sliding window, block 1000
118
# cannot be removed.
114
- removed = manager.remove_skipped_blocks(block_table, 4)
115
- assert_block_id(removed, [])
119
+ manager.remove_skipped_blocks("test", 4)
120
assert_block_id(block_table, original_block_ids)
121
122
# 5 tokens are computed. Token 0 & 1 are out of the sliding window.
123
# Block 1000 can be removed.
120
- removed = manager.remove_skipped_blocks(block_table, 5)
121
- assert_block_id(removed, [original_block_ids[0]])
124
+ manager.remove_skipped_blocks("test", 5)
125
assert_block_id(block_table, [null_block_id] + original_block_ids[1:])
126
127
# 6 tokens are computed. Token 0-2 are out of the sliding window.
128
# Cannot remove new block as the block 1001 is still used by token 3.
126
- removed = manager.remove_skipped_blocks(block_table, 6)
127
- assert_block_id(removed, [])
129
+ manager.remove_skipped_blocks("test", 6)
130
assert_block_id(block_table, [null_block_id] + original_block_ids[1:])
131
132
# 7 tokens are computed. Token 0-3 are out of the sliding window.
133
# Block 1001 can be removed and block 1000 is already removed.
132
- removed = manager.remove_skipped_blocks(block_table, 7)
133
- assert_block_id(removed, [original_block_ids[1]])
134
+ manager.remove_skipped_blocks("test", 7)
135
assert_block_id(block_table, [null_block_id] * 2 + original_block_ids[2:])
136
137
# 11 tokens are computed. Token 0-7 are out of the sliding window.
138
# Block 1002 & 1003 can be removed now. Block 1003 represents a longer
139
# sequence, and is expected to be evicted earlier than 1002, so the order
140
# of removed blocks should be [1003, 1002].
140
- removed = manager.remove_skipped_blocks(block_table, 11)
141
- assert_block_id(removed, [original_block_ids[3], original_block_ids[2]])
141
+ manager.remove_skipped_blocks("test", 11)
142
assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
vllm/v1/core/kv_cache_manager.py
+35
-100
@@ -1,17 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
from collections import defaultdict
4
-from collections.abc import Iterable
4
from dataclasses import dataclass
5
from typing import Optional
6
7
from vllm.distributed.kv_events import KVCacheEvent
8
from vllm.logger import init_logger
10
-from vllm.utils import cdiv, sha256
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,
12
hash_request_tokens)
14
-from vllm.v1.core.specialized_manager import get_specialized_manager
13
+from vllm.v1.core.specialized_manager import get_manager_for_kv_cache_spec
14
from vllm.v1.kv_cache_interface import KVCacheConfig
15
from vllm.v1.metrics.stats import PrefixCacheStats
16
from vllm.v1.request import Request, RequestStatus
@@ -56,7 +55,6 @@ class KVCacheManager:
55
self.block_size = kv_cache_spec.block_size
56
self.num_gpu_blocks = kv_cache_config.num_blocks
57
self.max_model_len = max_model_len
59
- self.max_num_blocks_per_req = cdiv(max_model_len, self.block_size)
58
59
self.enable_caching = enable_caching
60
self.caching_hash_fn = sha256 if caching_hash_algo == "sha256" else hash
@@ -68,30 +66,20 @@ class KVCacheManager:
66
self.block_pool = BlockPool(self.num_gpu_blocks, enable_caching,
67
enable_kv_cache_events)
68
71
- self.specialized_manager = get_specialized_manager(
69
+ self.single_type_manager = get_manager_for_kv_cache_spec(
70
kv_cache_spec=kv_cache_spec,
71
block_pool=self.block_pool,
72
use_eagle=self.use_eagle,
73
+ num_kv_cache_groups=1,
74
+ caching_hash_fn=self.caching_hash_fn,
75
)
76
77
- # Mapping from request ID to blocks to track the blocks allocated
78
- # for each request, so that we can free the blocks when the request
79
- # is finished.
80
- self.req_to_blocks: defaultdict[str,
81
- list[KVCacheBlock]] = defaultdict(list)
82
-
77
# Mapping from request ID to kv block hashes.
78
# This is to avoid recomputing the block hashes for each call of
79
# `get_computed_blocks` or `allocate_slots`.
80
self.req_to_block_hashes: defaultdict[
81
str, list[BlockHashType]] = defaultdict(list)
82
89
- # {req_id: The number of cached blocks for this given request}
90
- # This is used to track the number of cached blocks for each request.
91
- # This is only used to track the RUNNING requests, we do not track the
92
- # data for reempted ones.
93
- self.num_cached_block: dict[str, int] = {}
94
-
83
@property
84
def usage(self) -> float:
85
"""Get the KV cache usage.
@@ -159,7 +147,7 @@ class KVCacheManager:
147
last_block_hash = None
148
149
computed_blocks = (
162
- self.specialized_manager.find_longest_cache_hit(block_hashes))
150
+ self.single_type_manager.find_longest_cache_hit(block_hashes))
151
152
if self.log_stats:
153
assert self.prefix_cache_stats is not None
@@ -181,7 +169,7 @@ class KVCacheManager:
169
def allocate_slots(
170
self,
171
request: Request,
184
- num_tokens: int,
172
+ num_new_tokens: int,
173
new_computed_blocks: Optional[KVCacheBlocks] = None,
174
num_lookahead_tokens: int = 0,
175
) -> Optional[KVCacheBlocks]:
@@ -189,7 +177,7 @@ class KVCacheManager:
177
178
Args:
179
request: The request to allocate slots.
192
- num_tokens: The number of tokens to allocate, including external
180
+ num_new_tokens: The number of tokens to allocate, including external
181
tokens. Note that this does not include tokens that have
182
already been computed locally (i.e. new_computed_blocks).
183
new_computed_blocks: The new computed blocks just hitting the
@@ -215,44 +203,38 @@ class KVCacheManager:
203
Returns:
204
A list of new allocated blocks.
205
"""
218
- if num_tokens == 0:
219
- raise ValueError("num_tokens must be greater than 0")
206
+ if num_new_tokens == 0:
207
+ raise ValueError("num_new_tokens must be greater than 0")
208
209
if new_computed_blocks is not None:
210
new_computed_block_list = new_computed_blocks.blocks
211
else:
212
new_computed_block_list = []
213
226
- req_blocks = self.req_to_blocks[request.request_id]
227
-
214
# Free the blocks that are skipped during the attention computation
215
# (e.g., tokens outside the sliding window).
216
# We can do this even if we cannot schedule this request due to
217
# insufficient free blocks.
218
# Should call this function before allocating new blocks to reduce
219
# the number of evicted blocks.
234
- removed_blocks = self.specialized_manager.remove_skipped_blocks(
235
- req_blocks, request.num_computed_tokens)
236
- self.block_pool.free_blocks(removed_blocks)
220
+ self.single_type_manager.remove_skipped_blocks(
221
+ request.request_id, request.num_computed_tokens)
222
223
# The number of computed tokens is the number of computed tokens plus
224
# the new prefix caching hits
225
num_computed_tokens = (request.num_computed_tokens +
226
len(new_computed_block_list) * self.block_size)
242
- num_required_blocks = cdiv(
243
- num_computed_tokens + num_tokens + num_lookahead_tokens,
244
- self.block_size)
245
- num_new_blocks = (num_required_blocks - len(req_blocks) -
246
- len(new_computed_block_list))
247
-
248
- # If a computed block of a request is an eviction candidate (in the
249
- # free queue and ref_cnt == 0), it cannot be counted as a free block
250
- # when allocating this request.
251
- num_evictable_computed_blocks = sum(1
252
- for blk in new_computed_block_list
253
- if blk.ref_cnt == 0)
254
- if (num_new_blocks > self.block_pool.get_num_free_blocks() -
255
- num_evictable_computed_blocks):
227
+ num_tokens_need_slot = min(
228
+ num_computed_tokens + num_new_tokens + num_lookahead_tokens,
229
+ self.max_model_len)
230
+ num_blocks_to_allocate = (
231
+ self.single_type_manager.get_num_blocks_to_allocate(
232
+ request_id=request.request_id,
233
+ num_tokens=num_tokens_need_slot,
234
+ new_computed_blocks=new_computed_block_list,
235
+ ))
236
+
237
+ if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
238
# Cannot allocate new blocks
239
return None
240
@@ -266,74 +248,33 @@ class KVCacheManager:
248
249
# Append the new computed blocks to the request blocks until now to
250
# avoid the case where the new blocks cannot be allocated.
269
- req_blocks.extend(new_computed_block_list)
270
-
271
- # Start to handle new blocks
251
+ self.single_type_manager.save_new_computed_blocks(
252
+ request.request_id, new_computed_block_list)
253
273
- if num_new_blocks <= 0:
274
- # No new block is needed.
275
- new_blocks = []
276
- else:
277
- # Get new blocks from the free block pool.
278
- num_new_blocks = min(
279
- num_new_blocks,
280
- self.block_pool.get_num_free_blocks(),
281
- # Should not exceed the maximum number of blocks per request.
282
- # This is especially because the block table has the shape
283
- # [..., max_num_blocks_per_req].
284
- self.max_num_blocks_per_req - len(req_blocks),
285
- )
286
- assert num_new_blocks > 0
287
-
288
- # Concatenate the computed block IDs and the new block IDs.
289
- new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
290
- req_blocks.extend(new_blocks)
254
+ new_blocks = self.single_type_manager.allocate_new_blocks(
255
+ request.request_id, num_tokens_need_slot)
256
257
if not self.enable_caching:
258
return KVCacheBlocks(new_blocks)
259
295
- # Use `new_computed_block_list` for a new request, and
296
- # `num_cached_block` for a running request.
297
- num_cached_blocks = self.num_cached_block.get(
298
- request.request_id, len(new_computed_block_list))
260
# Speculated tokens might be rejected in the future, so we does
261
# not cache any speculated tokens. We only cache blocks with
262
# generated (accepted) tokens.
302
- num_full_blocks_after_append = (num_computed_tokens + num_tokens - len(
303
- request.spec_token_ids)) // self.block_size
304
-
305
- self.block_pool.cache_full_blocks(
306
- request=request,
307
- blocks=req_blocks,
308
- block_hashes=self.req_to_block_hashes[request.request_id],
309
- num_cached_blocks=num_cached_blocks,
310
- num_full_blocks=num_full_blocks_after_append,
311
- block_size=self.block_size,
312
- hash_fn=self.caching_hash_fn,
313
- )
263
+ self.single_type_manager.cache_blocks(
264
+ request, self.req_to_block_hashes[request.request_id],
265
+ num_computed_tokens + num_new_tokens - len(request.spec_token_ids))
266
315
- self.num_cached_block[
316
- request.request_id] = num_full_blocks_after_append
267
return KVCacheBlocks(new_blocks)
268
269
def free(self, request: Request) -> None:
270
"""Free the blocks allocated for the request.
321
- When caching is enabled, we free the blocks in reverse order so that
322
- the tail blocks are evicted first.
271
+ We free the blocks in reverse order so that he tail blocks are evicted
272
+ first when caching is enabled.
273
274
Args:
275
request: The request to free the blocks.
276
"""
327
- # Default to [] in case a request is freed (aborted) before alloc.
328
- blocks = self.req_to_blocks.pop(request.request_id, [])
329
- ordered_blocks: Iterable[KVCacheBlock] = blocks
330
- if self.enable_caching:
331
- # Free blocks in reverse order so that the tail blocks are
332
- # freed first.
333
- ordered_blocks = reversed(blocks)
334
-
335
- self.block_pool.free_blocks(ordered_blocks)
336
- self.num_cached_block.pop(request.request_id, None)
277
+ self.single_type_manager.free(request.request_id)
278
279
def reset_prefix_cache(self) -> bool:
280
"""Reset prefix cache. This function may be used in RLHF
@@ -390,14 +331,8 @@ class KVCacheManager:
331
int: The number of common prefix blocks.
332
"""
333
assert request.status == RequestStatus.RUNNING
393
- blocks = self.req_to_blocks[request.request_id]
394
- num_common_blocks = 0
395
- for block in blocks:
396
- if block.ref_cnt == num_running_requests:
397
- num_common_blocks += 1
398
- else:
399
- break
400
- return num_common_blocks
334
+ return self.single_type_manager.get_num_common_prefix_blocks(
335
+ request.request_id, num_running_requests)
336
337
def free_block_hashes(self, request: Request) -> None:
338
"""Discard the block hashes for the request.
vllm/v1/core/specialized_manager.py
+197
-24
@@ -1,17 +1,20 @@
1
# SPDX-License-Identifier: Apache-2.0
2
from abc import ABC, abstractmethod
3
+from collections import defaultdict
4
+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
9
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheSpec,
10
SlidingWindowSpec)
11
+from vllm.v1.request import Request
12
13
11
-class SpecializedManager(ABC):
14
+class SingleTypeKVCacheManager(ABC):
15
"""
13
- An abstract base class for specialized managers that handle the kv
14
- cache management logic of different attention layers.
16
+ An abstract base class for a manager that handle the kv cache management
17
+ logic of one specific type of attention layer.
18
"""
19
20
def __init__(
@@ -19,12 +22,18 @@ class SpecializedManager(ABC):
22
kv_cache_spec: KVCacheSpec,
23
block_pool: BlockPool,
24
use_eagle: bool,
25
+ num_kv_cache_groups: int,
26
+ caching_hash_fn: Callable,
27
) -> None:
28
"""
29
Initializes the SpecializedManager.
30
Args:
31
kv_cache_spec: The kv_cache_spec for this manager.
32
block_pool: The block pool.
33
+ use_eagle: Whether to use eagle.
34
+ num_kv_cache_groups: The number of kv cache groups managed by this
35
+ manager.
36
+ caching_hash_fn: The caching hash function.
37
"""
38
39
self.block_size = kv_cache_spec.block_size
@@ -34,6 +43,149 @@ class SpecializedManager(ABC):
43
# Needs special handling for find_longest_cache_hit if eagle is enabled
44
self.use_eagle = use_eagle
45
46
+ # Mapping from request ID to blocks to track the blocks allocated
47
+ # for each request, so that we can free the blocks when the request
48
+ # is finished.
49
+ self.req_to_blocks: defaultdict[str,
50
+ list[KVCacheBlock]] = defaultdict(list)
51
+
52
+ # {req_id: The number of cached blocks for this given request}
53
+ # This is used to track the number of cached blocks for each request.
54
+ # This is only used to track the RUNNING requests, we do not track the
55
+ # data for reempted ones.
56
+ self.num_cached_block: dict[str, int] = {}
57
+
58
+ self.num_kv_cache_groups = num_kv_cache_groups
59
+ self.caching_hash_fn = caching_hash_fn
60
+
61
+ def get_num_blocks_to_allocate(
62
+ self, request_id: str, num_tokens: int,
63
+ new_computed_blocks: list[KVCacheBlock]) -> int:
64
+ """
65
+ Get the number of blocks needed to be allocated for the request.
66
+
67
+ Args:
68
+ request_id: The request ID.
69
+ num_tokens: The total number of tokens that need a slot (including
70
+ tokens that are already allocated).
71
+ new_computed_blocks: The new computed blocks just hitting the
72
+ prefix caching.
73
+
74
+ Returns:
75
+ The number of blocks.
76
+ """
77
+
78
+ num_required_blocks = cdiv(num_tokens, self.block_size)
79
+ num_new_blocks = (num_required_blocks - len(new_computed_blocks) -
80
+ len(self.req_to_blocks[request_id]))
81
+ # If a computed block of a request is an eviction candidate (in the
82
+ # free queue and ref_cnt == 0), it will be changed from a free block
83
+ # to a computed block when the request is allocated, so we also count
84
+ # it as needed to be allocated.
85
+ num_evictable_computed_blocks = sum(blk.ref_cnt == 0
86
+ for blk in new_computed_blocks)
87
+ return ((num_new_blocks + num_evictable_computed_blocks) *
88
+ self.num_kv_cache_groups)
89
+
90
+ def save_new_computed_blocks(
91
+ self, request_id: str,
92
+ new_computed_blocks: list[KVCacheBlock]) -> None:
93
+ """
94
+ Add the new computed blocks to the request.
95
+
96
+ Args:
97
+ request_id: The request ID.
98
+ new_computed_blocks: The new computed blocks just hitting the
99
+ prefix cache.
100
+ """
101
+ if request_id not in self.num_cached_block:
102
+ # A new request.
103
+ req_blocks = self.req_to_blocks[request_id]
104
+ assert len(req_blocks) == 0
105
+ req_blocks.extend(new_computed_blocks)
106
+ self.num_cached_block[request_id] = len(new_computed_blocks)
107
+ else:
108
+ # A running request. Should not have new computed blocks.
109
+ assert len(new_computed_blocks) == 0
110
+
111
+ def allocate_new_blocks(self, request_id: str,
112
+ num_tokens: int) -> list[KVCacheBlock]:
113
+ """
114
+ Allocate new blocks for the request to give it at least `num_tokens`
115
+ token slots.
116
+
117
+ Args:
118
+ request_id: The request ID.
119
+ num_tokens: The total number of tokens that need a slot (including
120
+ tokens that are already allocated).
121
+
122
+ Returns:
123
+ The new allocated blocks.
124
+ """
125
+ req_blocks = self.req_to_blocks[request_id]
126
+ num_required_blocks = cdiv(num_tokens, self.block_size)
127
+ num_new_blocks = num_required_blocks - len(req_blocks)
128
+ if num_new_blocks <= 0:
129
+ return []
130
+ else:
131
+ new_blocks = self.block_pool.get_new_blocks(
132
+ num_new_blocks * self.num_kv_cache_groups)
133
+ req_blocks.extend(new_blocks)
134
+ return new_blocks
135
+
136
+ def cache_blocks(self, request: Request, block_hashes: list[BlockHashType],
137
+ num_tokens: int) -> None:
138
+ """
139
+ Cache the blocks for the request.
140
+
141
+ Args:
142
+ request: The request.
143
+ block_hashes: The block hashes of the request.
144
+ num_tokens: The total number of tokens that need to be cached
145
+ (including tokens that are already cached).
146
+ """
147
+ num_cached_blocks = self.num_cached_block[request.request_id]
148
+ num_full_blocks = num_tokens // self.block_size
149
+
150
+ self.block_pool.cache_full_blocks(
151
+ request=request,
152
+ blocks=self.req_to_blocks[request.request_id],
153
+ block_hashes=block_hashes,
154
+ num_cached_blocks=num_cached_blocks,
155
+ num_full_blocks=num_full_blocks,
156
+ block_size=self.block_size,
157
+ hash_fn=self.caching_hash_fn,
158
+ )
159
+
160
+ self.num_cached_block[request.request_id] = num_full_blocks
161
+
162
+ def free(self, request_id: str) -> None:
163
+ # Default to [] in case a request is freed (aborted) before alloc.
164
+ req_blocks = self.req_to_blocks.pop(request_id, [])
165
+
166
+ # Free blocks in reverse order so that the tail blocks are
167
+ # freed first.
168
+ ordered_blocks = reversed(req_blocks)
169
+
170
+ self.block_pool.free_blocks(ordered_blocks)
171
+ self.num_cached_block.pop(request_id, None)
172
+
173
+ @abstractmethod
174
+ def get_num_common_prefix_blocks(self, request_id: str,
175
+ num_running_requests: int) -> int:
176
+ """
177
+ Get the number of common prefix blocks for a request.
178
+
179
+ Args:
180
+ request_id: The request ID.
181
+ block_hashes: The block hashes of the request.
182
+
183
+ Returns:
184
+ The number of common prefix blocks.
185
+ """
186
+
187
+ raise NotImplementedError
188
+
189
@abstractmethod
190
def find_longest_cache_hit(
191
self, block_hashes: list[BlockHashType]) -> list[KVCacheBlock]:
@@ -41,7 +193,8 @@ class SpecializedManager(ABC):
193
Get the longest cache hit prefix of the blocks. If no cache hit is
194
found, return an empty list. if eagle is enabled, drop the last matched
195
block to force recompute the last block to get the required hidden
44
- states for eagle drafting head.
196
+ states for eagle drafting head. Need to be customized for each attention
197
+ type.
198
199
Args:
200
block_hashes: The block hashes of the request.
@@ -55,24 +208,23 @@ class SpecializedManager(ABC):
208
raise NotImplementedError
209
210
@abstractmethod
58
- def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
59
- num_computed_tokens: int) -> list[KVCacheBlock]:
211
+ def remove_skipped_blocks(self, request_id: str,
212
+ num_computed_tokens: int) -> None:
213
"""
214
Remove the blocks that are no longer needed from `blocks`. The removed
215
blocks should be replaced by null_block. Return the removed blocks in
216
eviction order, where the first returned block should be evicted first.
64
- Don't free the removed blocks in this function.
217
+ Don't free the removed blocks in this function. Need to be customized
218
+ for each attention type.
219
220
Args:
67
- blocks: The list of blocks to be updated.
221
+ request_id: The request ID.
222
num_computed_tokens: The number of tokens that have been computed.
69
- Returns:
70
- The removed blocks in eviction order.
223
"""
224
raise NotImplementedError
225
226
75
-class FullAttentionManager(SpecializedManager):
227
+class FullAttentionManager(SingleTypeKVCacheManager):
228
229
def find_longest_cache_hit(
230
self, block_hashes: list[BlockHashType]) -> list[KVCacheBlock]:
@@ -89,17 +241,28 @@ class FullAttentionManager(SpecializedManager):
241
computed_blocks.pop()
242
return computed_blocks
243
92
- def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
93
- num_computed_tokens: int) -> list[KVCacheBlock]:
244
+ def remove_skipped_blocks(self, request_id: str,
245
+ num_computed_tokens: int) -> None:
246
# No need to remove blocks for full attention.
95
- return []
247
+ pass
248
+
249
+ def get_num_common_prefix_blocks(self, request_id: str,
250
+ num_running_requests: int) -> int:
251
+ blocks = self.req_to_blocks[request_id]
252
+ num_common_blocks = 0
253
+ for block in blocks:
254
+ if block.ref_cnt == num_running_requests:
255
+ num_common_blocks += 1
256
+ else:
257
+ break
258
+ return num_common_blocks
259
260
98
-class SlidingWindowManager(SpecializedManager):
261
+class SlidingWindowManager(SingleTypeKVCacheManager):
262
263
def __init__(self, kv_cache_spec: SlidingWindowSpec, block_pool: BlockPool,
101
- use_eagle: bool):
102
- super().__init__(kv_cache_spec, block_pool, use_eagle)
264
+ use_eagle: bool, **kwargs) -> None:
265
+ super().__init__(kv_cache_spec, block_pool, use_eagle, **kwargs)
266
self.sliding_window = kv_cache_spec.sliding_window
267
# The number of contiguous blocks needed for prefix cache hit.
268
# -1 since the input token itself is also included in the window
@@ -148,13 +311,13 @@ class SlidingWindowManager(SpecializedManager):
311
computed_blocks.pop()
312
return computed_blocks
313
151
- def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
152
- num_computed_tokens: int) -> list[KVCacheBlock]:
314
+ def remove_skipped_blocks(self, request_id: str,
315
+ num_computed_tokens: int) -> None:
316
# Remove the blocks that are no longer be in the sliding window and
317
# skipped during the attention computation.
318
last_useful_token = num_computed_tokens - self.sliding_window + 1
319
last_useful_block = last_useful_token // self.block_size
157
-
320
+ blocks = self.req_to_blocks[request_id]
321
removed_blocks: list[KVCacheBlock] = []
322
for i in range(last_useful_block - 1, -1, -1):
323
if blocks[i] == self._null_block:
@@ -164,17 +327,27 @@ class SlidingWindowManager(SpecializedManager):
327
break
328
removed_blocks.append(blocks[i])
329
blocks[i] = self._null_block
167
- return removed_blocks
330
+ self.block_pool.free_blocks(removed_blocks)
331
+
332
+ def get_num_common_prefix_blocks(self, request_id: str,
333
+ num_running_requests: int) -> int:
334
+ """
335
+ NOTE(Chen): The prefix blocks are null blocks for sliding window layers.
336
+ So it's not correct to count ref_cnt like FullAttentionManager. Return
337
+ 0 here for correctness. Need to support cascade attention + sliding
338
+ window in the future.
339
+ """
340
+ return 0
341
342
170
-spec_manager_map: dict[type[KVCacheSpec], type[SpecializedManager]] = {
343
+spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
344
FullAttentionSpec: FullAttentionManager,
345
SlidingWindowSpec: SlidingWindowManager,
346
}
347
348
176
-def get_specialized_manager(kv_cache_spec: KVCacheSpec,
177
- **kwargs) -> SpecializedManager:
349
+def get_manager_for_kv_cache_spec(kv_cache_spec: KVCacheSpec,
350
+ **kwargs) -> SingleTypeKVCacheManager:
351
manager_class = spec_manager_map[type(kv_cache_spec)]
352
manager = manager_class(kv_cache_spec, **kwargs)
353
return manager