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, 16:00 UTC
[V1] Implement sliding window attention in kv_cache_manager (#14097)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Apr 1, 2025 at 15:33 UTC
3a5f0afcd2e91264f9ac812badaa0ed3cd0a623e
15 files changed
+662
-158
tests/core/block/e2e/test_correctness_sliding_window.py
+11
-4
@@ -129,12 +129,16 @@ def test_sliding_window_chunked_prefill(test_llm_generator, batch_size, seed,
129
check_answers(indices, answer, test_texts)
130
131
132
-def prep_prompts(batch_size: int):
132
+def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
133
"""
134
Generate prompts which a bunch of assignments,
135
then asking for the value of one of them.
136
The prompt is just under 10k tokens; sliding window is 4k
137
so the answer is outside sliding window, but should still be correct.
138
+
139
+ Args:
140
+ batch_size: number of prompts to generate
141
+ ln_range: an argument to control the length of the prompt
142
"""
143
prompts: list[str] = []
144
answer: list[int] = []
@@ -145,7 +149,7 @@ def prep_prompts(batch_size: int):
149
indices.append(idx)
150
prompt = "```python\n# We set a number of variables, " + \
151
f"x{idx} will be important later\n"
148
- ln = random.randint(800, 1100)
152
+ ln = random.randint(*ln_range)
153
for k in range(30, ln):
154
v = random.randint(10, 99)
155
if k == idx:
@@ -157,7 +161,10 @@ def prep_prompts(batch_size: int):
161
return prompts, answer, indices
162
163
160
-def check_answers(indices: list[int], answer: list[int], outputs: list[str]):
164
+def check_answers(indices: list[int],
165
+ answer: list[int],
166
+ outputs: list[str],
167
+ accept_rate: float = 0.7):
168
answer2 = [int(text[0:2].strip()) for text in outputs]
169
print(list(zip(indices, zip(answer, answer2))))
170
numok = 0
@@ -166,7 +173,7 @@ def check_answers(indices: list[int], answer: list[int], outputs: list[str]):
173
numok += 1
174
frac_ok = numok / len(answer)
175
print(f"Num OK: {numok}/{len(answer)} {frac_ok}")
169
- assert frac_ok > 0.7
176
+ assert frac_ok >= accept_rate
177
178
179
def check_window(prompts: list[str]):
tests/v1/core/test_prefix_caching.py
+63
-70
@@ -4,6 +4,7 @@
4
from typing import Optional
5
6
import pytest
7
+import torch
8
9
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
10
from vllm.sampling_params import SamplingParams
@@ -12,6 +13,8 @@ from vllm.v1.core.block_pool import BlockPool
13
from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
14
from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
15
hash_block_tokens)
16
+from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
17
+ KVCacheGroupSpec)
18
19
20
def make_request(request_id,
@@ -39,13 +42,23 @@ def make_request(request_id,
42
)
43
44
45
+def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
46
+ return KVCacheConfig(
47
+ num_blocks=num_blocks,
48
+ tensors={},
49
+ kv_cache_groups=[
50
+ KVCacheGroupSpec(['layer'],
51
+ FullAttentionSpec(block_size, 1, 1, torch.float32,
52
+ False))
53
+ ],
54
+ )
55
+
56
+
57
@pytest.mark.parametrize("hash_algo", ["sha256", "hash"])
58
def test_prefill(hash_algo):
59
manager = KVCacheManager(
45
- block_size=16,
46
- num_gpu_blocks=10,
60
+ make_kv_cache_config(16, 11),
61
max_model_len=8192,
48
- sliding_window=None,
62
enable_caching=True,
63
caching_hash_algo=hash_algo,
64
num_preallocate_tokens=16,
@@ -67,12 +80,12 @@ def test_prefill(hash_algo):
80
assert not computed_blocks
81
assert num_computed_tokens == 0
82
blocks = manager.allocate_slots(req0, 55, computed_blocks)
70
- assert [b.block_id for b in blocks] == [0, 1, 2, 3, 4]
83
+ assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]
84
85
# Check full block metadata
86
parent_block_hash = None
74
- for block_id in (0, 1, 2):
75
- block_tokens = tuple(all_token_ids[block_id * 16:(block_id + 1) * 16])
87
+ for block_id in (1, 2, 3):
88
+ block_tokens = tuple(all_token_ids[(block_id - 1) * 16:block_id * 16])
89
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
90
block_tokens)
91
assert manager.block_pool.blocks[block_id].block_hash == block_hash
@@ -80,7 +93,7 @@ def test_prefill(hash_algo):
93
parent_block_hash = block_hash.hash_value
94
95
# Check partial/preallocated block metadata
83
- for block_id in (3, 4):
96
+ for block_id in (4, 5):
97
assert manager.block_pool.blocks[block_id].block_hash is None
98
assert manager.block_pool.blocks[block_id].ref_cnt == 1
99
@@ -90,11 +103,11 @@ def test_prefill(hash_algo):
103
req1 = make_request("1", common_token_ids + unique_token_ids)
104
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
105
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
93
- assert [b.block_id for b in computed_blocks] == [0, 1, 2]
106
+ assert [b.block_id for b in computed_blocks] == [1, 2, 3]
107
assert num_computed_tokens == 3 * 16
108
num_new_tokens = 53 - 3 * 16
109
blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)
97
- assert [b.block_id for b in blocks] == [5, 6]
110
+ assert [b.block_id for b in blocks] == [6, 7]
111
for block in computed_blocks:
112
assert block.ref_cnt == 2
113
@@ -107,14 +120,14 @@ def test_prefill(hash_algo):
120
# All blocks should be available.
121
assert manager.block_pool.free_block_queue.num_free_blocks == 10
122
# The order should be
110
- # [unallocated (7, 8, 9)]
111
- # [unique_req0 (4, 3)]
112
- # [unique_req1 (6, 5)]
113
- # [common (2, 1, 0)]
123
+ # [unallocated (8, 9, 10)]
124
+ # [unique_req0 (5, 4)]
125
+ # [unique_req1 (7, 6)]
126
+ # [common (3, 2, 1)]
127
assert [
128
b.block_id
129
for b in manager.block_pool.free_block_queue.get_all_free_blocks()
117
- ] == [7, 8, 9, 4, 3, 6, 5, 2, 1, 0]
130
+ ] == [8, 9, 10, 5, 4, 7, 6, 3, 2, 1]
131
132
# Cache hit in the common prefix when the original block is already free.
133
# Incomplete 1 block (6 tokens)
@@ -122,11 +135,11 @@ def test_prefill(hash_algo):
135
req2 = make_request("2", common_token_ids + unique_token_ids)
136
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
137
assert len(manager.req_to_block_hashes[req2.request_id]) == 3
125
- assert [b.block_id for b in computed_blocks] == [0, 1, 2]
138
+ assert [b.block_id for b in computed_blocks] == [1, 2, 3]
139
assert num_computed_tokens == 3 * 16
140
num_new_tokens = 53 - 3 * 16
141
blocks = manager.allocate_slots(req2, num_new_tokens, computed_blocks)
129
- assert [b.block_id for b in blocks] == [7, 8]
142
+ assert [b.block_id for b in blocks] == [8, 9]
143
144
# Although we only have 5 free blocks, we have 8 blocks in
145
# the free block queue due to lazy removal.
@@ -148,7 +161,7 @@ def test_prefill(hash_algo):
161
assert num_computed_tokens == 0
162
blocks = manager.allocate_slots(req3, 16 * 9, computed_blocks)
163
# This block ID order also checks the eviction order.
151
- assert [b.block_id for b in blocks] == [9, 4, 3, 6, 5, 8, 7, 2, 1, 0]
164
+ assert [b.block_id for b in blocks] == [10, 5, 4, 7, 6, 9, 8, 3, 2, 1]
165
assert manager.block_pool.free_block_queue.num_free_blocks == 0
166
assert manager.block_pool.free_block_queue.free_list_head is None
167
assert manager.block_pool.free_block_queue.free_list_tail is None
@@ -162,10 +175,8 @@ def test_prefill_plp():
175
3. Schedule plp request; no hit should occur; validate blocks
176
'''
177
manager = KVCacheManager(
165
- block_size=16,
166
- num_gpu_blocks=10,
178
+ make_kv_cache_config(16, 11),
179
max_model_len=8192,
168
- sliding_window=None,
180
enable_caching=True,
181
num_preallocate_tokens=16,
182
)
@@ -186,13 +197,13 @@ def test_prefill_plp():
197
assert not computed_blocks
198
assert num_computed_tokens == 0
199
blocks = manager.allocate_slots(req0, 55, computed_blocks)
189
- assert [b.block_id for b in blocks] == [0, 1, 2, 3, 4]
200
+ assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]
201
req0_block_hashes = [b.block_hash for b in blocks]
202
203
# Check full block metadata
204
parent_block_hash = None
194
- for block_id in (0, 1, 2):
195
- block_tokens = tuple(all_token_ids[block_id * 16:(block_id + 1) * 16])
205
+ for block_id in (1, 2, 3):
206
+ block_tokens = tuple(all_token_ids[(block_id - 1) * 16:block_id * 16])
207
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
208
block_tokens)
209
assert manager.block_pool.blocks[block_id].block_hash == block_hash
@@ -200,7 +211,7 @@ def test_prefill_plp():
211
parent_block_hash = block_hash.hash_value
212
213
# Check partial/preallocated block metadata
203
- for block_id in (3, 4):
214
+ for block_id in (4, 5):
215
assert manager.block_pool.blocks[block_id].block_hash is None
216
assert manager.block_pool.blocks[block_id].ref_cnt == 1
217
@@ -211,11 +222,11 @@ def test_prefill_plp():
222
req1 = make_request("1", common_token_ids + unique_token_ids)
223
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
224
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
214
- assert [b.block_id for b in computed_blocks] == [0, 1, 2]
225
+ assert [b.block_id for b in computed_blocks] == [1, 2, 3]
226
assert num_computed_tokens == 3 * 16
227
num_new_tokens = 53 - 3 * 16
228
blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)
218
- assert [b.block_id for b in blocks] == [5, 6]
229
+ assert [b.block_id for b in blocks] == [6, 7]
230
for block in computed_blocks:
231
assert block.ref_cnt == 2
232
@@ -228,14 +239,14 @@ def test_prefill_plp():
239
# All blocks should be available.
240
assert manager.block_pool.free_block_queue.num_free_blocks == 10
241
# The order should be
231
- # [unallocated (7, 8, 9)]
232
- # [unique_req0 (4, 3)]
233
- # [unique_req1 (6, 5)]
234
- # [common (2, 1, 0)]
242
+ # [unallocated (8, 9, 10)]
243
+ # [unique_req0 (5, 4)]
244
+ # [unique_req1 (7, 6)]
245
+ # [common (3, 2, 1)]
246
assert [
247
b.block_id
248
for b in manager.block_pool.free_block_queue.get_all_free_blocks()
238
- ] == [7, 8, 9, 4, 3, 6, 5, 2, 1, 0]
249
+ ] == [8, 9, 10, 5, 4, 7, 6, 3, 2, 1]
250
251
# Request #2 is a prompt-logprobs request:
252
# NO cache hit in the common prefix; duplicates request #0 cached blocks
@@ -251,7 +262,7 @@ def test_prefill_plp():
262
block_ids = [b.block_id for b in blocks]
263
# Duplicate cached blocks have different ids but same hashes vs request #0
264
assert [b.block_hash for b in blocks] == req0_block_hashes
254
- assert block_ids != [0, 1, 2, 3, 4]
265
+ assert block_ids != [1, 2, 3, 4, 5]
266
267
# Request #2 block hashes are valid since request #0 hashes are.
268
# Check block reference counts.
@@ -263,10 +274,8 @@ def test_prefill_plp():
274
275
def test_decode():
276
manager = KVCacheManager(
266
- block_size=16,
267
- num_gpu_blocks=10,
277
+ make_kv_cache_config(16, 11),
278
max_model_len=8192,
269
- sliding_window=None,
279
enable_caching=True,
280
num_preallocate_tokens=16,
281
)
@@ -282,7 +291,7 @@ def test_decode():
291
assert not computed_blocks
292
assert num_computed_tokens == 0
293
blocks = manager.allocate_slots(req0, 55, computed_blocks)
285
- assert [b.block_id for b in blocks] == [0, 1, 2, 3, 4]
294
+ assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]
295
296
# Append slots without allocating a new block.
297
req0.num_computed_tokens = 55
@@ -316,10 +325,8 @@ def test_decode():
325
326
def test_evict():
327
manager = KVCacheManager(
319
- block_size=16,
320
- num_gpu_blocks=10,
328
+ make_kv_cache_config(16, 11),
329
max_model_len=8192,
322
- sliding_window=None,
330
enable_caching=True,
331
num_preallocate_tokens=16,
332
)
@@ -350,15 +357,15 @@ def test_evict():
357
assert [
358
b.block_id
359
for b in manager.block_pool.free_block_queue.get_all_free_blocks()
353
- ] == [6, 5, 4, 3, 2, 1, 0, 9, 8, 7]
360
+ ] == [7, 6, 5, 4, 3, 2, 1, 10, 9, 8]
361
362
# Touch the first 2 blocks.
363
req2 = make_request("2", list(range(2 * 16 + 3)))
364
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
358
- assert [b.block_id for b in computed_blocks] == [0, 1]
365
+ assert [b.block_id for b in computed_blocks] == [1, 2]
366
assert num_computed_tokens == 2 * 16
367
blocks = manager.allocate_slots(req2, 3, computed_blocks)
361
- assert [b.block_id for b in blocks] == [6, 5]
368
+ assert [b.block_id for b in blocks] == [7, 6]
369
assert manager.block_pool.free_block_queue.num_free_blocks == 6
370
371
@@ -369,10 +376,8 @@ def test_hash_block_correct_reuse():
376
"""
377
block_size = 16
378
manager = KVCacheManager(
372
- block_size=block_size,
373
- num_gpu_blocks=1,
379
+ make_kv_cache_config(16, 2),
380
max_model_len=8192,
375
- sliding_window=None,
381
enable_caching=True,
382
num_preallocate_tokens=0,
383
)
@@ -408,10 +413,8 @@ def test_computed_blocks_not_evicted():
413
"""
414
block_size = 16
415
manager = KVCacheManager(
411
- block_size=block_size,
412
- num_gpu_blocks=2,
416
+ make_kv_cache_config(block_size, 3),
417
max_model_len=8192,
414
- sliding_window=None,
418
enable_caching=True,
419
num_preallocate_tokens=0,
420
)
@@ -424,7 +427,7 @@ def test_computed_blocks_not_evicted():
427
assert num_computed_tokens == 0
428
blocks = manager.allocate_slots(req0, num_tokens, computed_blocks)
429
assert len(blocks) == 1
427
- assert blocks[0].block_id == 0
430
+ assert blocks[0].block_id == 1
431
432
# Allocate another block.
433
req1 = make_request("1", list(range(num_tokens, num_tokens * 2)))
@@ -433,7 +436,7 @@ def test_computed_blocks_not_evicted():
436
assert num_computed_tokens == 0
437
blocks = manager.allocate_slots(req1, num_tokens, computed_blocks)
438
assert len(blocks) == 1
436
- assert blocks[0].block_id == 1
439
+ assert blocks[0].block_id == 2
440
441
# Free the blocks.
442
manager.free(req0)
@@ -444,13 +447,13 @@ def test_computed_blocks_not_evicted():
447
req2 = make_request("2", list(range(num_tokens * 2)))
448
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
449
assert len(computed_blocks) == 1
447
- assert computed_blocks[0].block_id == 0
450
+ assert computed_blocks[0].block_id == 1
451
assert num_computed_tokens == block_size
452
453
blocks = manager.allocate_slots(req2, num_tokens * 2 - num_tokens,
454
computed_blocks)
455
assert len(blocks) == 1
453
- assert blocks[0].block_id == 1
456
+ assert blocks[0].block_id == 2
457
458
459
def test_basic_prefix_caching_disabled():
@@ -459,10 +462,8 @@ def test_basic_prefix_caching_disabled():
462
"""
463
block_size = 4
464
manager = KVCacheManager(
462
- block_size=block_size,
463
- num_gpu_blocks=4,
465
+ make_kv_cache_config(block_size, 5),
466
max_model_len=8192,
465
- sliding_window=None,
467
enable_caching=False,
468
num_preallocate_tokens=0,
469
)
@@ -502,10 +503,8 @@ def test_preallocate_blocks(num_preallocate_tokens: int, block_size: int):
503
This tests that the preallocated blocks are correctly added.
504
"""
505
manager = KVCacheManager(
505
- block_size=block_size,
506
- num_gpu_blocks=10,
506
+ make_kv_cache_config(block_size, 11),
507
max_model_len=8192,
508
- sliding_window=None,
508
enable_caching=True,
509
num_preallocate_tokens=num_preallocate_tokens,
510
)
@@ -586,10 +585,8 @@ def test_mm_prefix_caching():
585
This tests that the multi-modal prefix caching is correct.
586
"""
587
manager = KVCacheManager(
589
- block_size=16,
590
- num_gpu_blocks=10,
588
+ make_kv_cache_config(16, 11),
589
max_model_len=8192,
592
- sliding_window=None,
590
enable_caching=True,
591
num_preallocate_tokens=16,
592
)
@@ -629,7 +626,7 @@ def test_mm_prefix_caching():
626
assert block_hashes[2].extra_keys == ("bbb", )
627
628
blocks = manager.allocate_slots(req0, 59, computed_blocks)
632
- assert [b.block_id for b in blocks] == [0, 1, 2, 3, 4]
629
+ assert [b.block_id for b in blocks] == [1, 2, 3, 4, 5]
630
req0.num_computed_tokens = 59
631
632
# Append slots without allocating a new block.
@@ -667,10 +664,8 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
664
"""
665
block_size = 16
666
manager = KVCacheManager(
670
- block_size=block_size,
671
- num_gpu_blocks=10,
667
+ make_kv_cache_config(block_size, 11),
668
max_model_len=8192,
673
- sliding_window=None,
669
enable_caching=True,
670
num_preallocate_tokens=0,
671
)
@@ -723,10 +718,8 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
718
719
def test_reset_prefix_cache():
720
manager = KVCacheManager(
726
- block_size=16,
727
- num_gpu_blocks=10,
721
+ make_kv_cache_config(16, 11),
722
max_model_len=8192,
729
- sliding_window=None,
723
enable_caching=True,
724
num_preallocate_tokens=0,
725
)
@@ -736,7 +729,7 @@ def test_reset_prefix_cache():
729
all_token_ids = full_block_token_ids + unique_token_ids
730
req0 = make_request("0", all_token_ids)
731
blocks = manager.allocate_slots(req0, 55)
739
- assert [b.block_id for b in blocks] == [0, 1, 2, 3]
732
+ assert [b.block_id for b in blocks] == [1, 2, 3, 4]
733
734
unique_token_ids = [4] * 7
735
all_token_ids = full_block_token_ids + unique_token_ids
@@ -745,7 +738,7 @@ def test_reset_prefix_cache():
738
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
739
assert len(computed_blocks) == 3
740
blocks = manager.allocate_slots(req1, 7, computed_blocks)
748
- assert [b.block_id for b in blocks] == [4]
741
+ assert [b.block_id for b in blocks] == [5]
742
743
# Failed to reset prefix cache because some blocks are not freed yet.
744
assert not manager.reset_prefix_cache()
tests/v1/core/test_scheduler.py
+12
@@ -2,12 +2,15 @@
2
from typing import Optional
3
4
import pytest
5
+import torch
6
7
from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
8
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
9
from vllm.sampling_params import SamplingParams
10
from vllm.v1.core.sched.output import SchedulerOutput
11
from vllm.v1.core.sched.scheduler import Scheduler
12
+from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
13
+ KVCacheGroupSpec)
14
from vllm.v1.outputs import ModelRunnerOutput
15
from vllm.v1.request import Request, RequestStatus
16
from vllm.v1.structured_output import StructuredOutputManager
@@ -66,12 +69,21 @@ def create_scheduler(
69
model_config=model_config,
70
cache_config=cache_config,
71
)
72
+ kv_cache_config = KVCacheConfig(
73
+ num_blocks=10000, # A large number of blocks to hold all requests
74
+ tensors={},
75
+ kv_cache_groups=[
76
+ KVCacheGroupSpec(['layer'],
77
+ FullAttentionSpec(16, 1, 1, torch.float32, False))
78
+ ],
79
+ )
80
cache_config.num_gpu_blocks = 10000
81
return Scheduler(
82
scheduler_config,
83
model_config,
84
cache_config,
85
lora_config=None,
86
+ kv_cache_config=kv_cache_config,
87
log_stats=True,
88
structured_output_manager=StructuredOutputManager(vllm_config),
89
)
tests/v1/core/test_specialized_manager.py
new
+138
@@ -0,0 +1,138 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import torch
4
+
5
+from vllm.v1.core.block_pool import BlockPool
6
+from vllm.v1.core.kv_cache_utils import BlockHashType, KVCacheBlock
7
+from vllm.v1.core.specialized_manager import SlidingWindowManager
8
+from vllm.v1.kv_cache_interface import SlidingWindowSpec
9
+
10
+
11
+def test_sliding_window_possible_cached_prefix():
12
+ sliding_window_spec = SlidingWindowSpec(
13
+ block_size=2,
14
+ num_kv_heads=1,
15
+ head_size=1,
16
+ dtype=torch.float32,
17
+ sliding_window=4,
18
+ use_mla=False,
19
+ )
20
+
21
+ block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
22
+ manager = SlidingWindowManager(sliding_window_spec, block_pool)
23
+
24
+ def run_one_case(block_is_cached, expect_length):
25
+ block_hash_list = [
26
+ BlockHashType(i, ()) for i in range(len(block_is_cached))
27
+ ]
28
+
29
+ block_pool.cached_block_hash_to_block.clear()
30
+
31
+ # Mock the block pool with the cached blocks
32
+ for i, (block_hash,
33
+ is_cached) in enumerate(zip(block_hash_list, block_is_cached)):
34
+ if is_cached:
35
+ block_pool.cached_block_hash_to_block[block_hash] = {
36
+ i: block_pool.blocks[i + 10]
37
+ }
38
+
39
+ computed_blocks = manager.find_longest_cache_hit(block_hash_list)
40
+ assert len(computed_blocks) == expect_length
41
+
42
+ assert all(block == block_pool.null_block
43
+ for block in computed_blocks[:expect_length - 2])
44
+ for i in range(2):
45
+ if i < expect_length:
46
+ block_index = expect_length - i - 1
47
+ assert computed_blocks[
48
+ block_index].block_id == block_index + 10
49
+
50
+ run_one_case([False] * 10, 0)
51
+ run_one_case([True], 1)
52
+ run_one_case([True, False], 1)
53
+ run_one_case([True, True], 2)
54
+ run_one_case([True, True, False], 2)
55
+ run_one_case([True, True, True], 3)
56
+ run_one_case([True, True, True, False], 3)
57
+ run_one_case([
58
+ True, True, False, True, False, False, True, True, False, True, True,
59
+ True
60
+ ], 12)
61
+ run_one_case([
62
+ True, True, False, True, False, False, True, True, False, False, False
63
+ ], 8)
64
+ run_one_case([
65
+ True, True, False, True, False, False, True, True, False, False, False,
66
+ True
67
+ ], 8)
68
+
69
+
70
+def test_sliding_window_remove_skipped_blocks():
71
+ sliding_window_spec = SlidingWindowSpec(
72
+ block_size=2,
73
+ num_kv_heads=1,
74
+ head_size=1,
75
+ dtype=torch.float32,
76
+ sliding_window=4,
77
+ use_mla=False,
78
+ )
79
+
80
+ block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
81
+
82
+ manager = SlidingWindowManager(sliding_window_spec, block_pool)
83
+
84
+ null_block_id = block_pool.null_block.block_id
85
+
86
+ def id_to_block_table(ids):
87
+ return [
88
+ KVCacheBlock(id_)
89
+ if id_ != null_block_id else block_pool.null_block for id_ in ids
90
+ ]
91
+
92
+ def assert_block_id(block_table, ids):
93
+ for block, id_ in zip(block_table, ids):
94
+ if id_ == null_block_id:
95
+ assert block == block_pool.null_block
96
+ else:
97
+ assert block.block_id == id_
98
+
99
+ original_block_ids = [
100
+ 1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010
101
+ ]
102
+ block_table = id_to_block_table(original_block_ids)
103
+ removed = manager.remove_skipped_blocks(block_table, 0)
104
+ assert_block_id(removed, [])
105
+ assert_block_id(block_table, original_block_ids)
106
+
107
+ # 4 tokens are computed. Only token 0 is out of the sliding window. As
108
+ # block 1000 also contains token 1 that is in the sliding window, block 1000
109
+ # cannot be removed.
110
+ removed = manager.remove_skipped_blocks(block_table, 4)
111
+ assert_block_id(removed, [])
112
+ assert_block_id(block_table, original_block_ids)
113
+
114
+ # 5 tokens are computed. Token 0 & 1 are out of the sliding window.
115
+ # Block 1000 can be removed.
116
+ removed = manager.remove_skipped_blocks(block_table, 5)
117
+ assert_block_id(removed, [original_block_ids[0]])
118
+ assert_block_id(block_table, [null_block_id] + original_block_ids[1:])
119
+
120
+ # 6 tokens are computed. Token 0-2 are out of the sliding window.
121
+ # Cannot remove new block as the block 1001 is still used by token 3.
122
+ removed = manager.remove_skipped_blocks(block_table, 6)
123
+ assert_block_id(removed, [])
124
+ assert_block_id(block_table, [null_block_id] + original_block_ids[1:])
125
+
126
+ # 7 tokens are computed. Token 0-3 are out of the sliding window.
127
+ # Block 1001 can be removed and block 1000 is already removed.
128
+ removed = manager.remove_skipped_blocks(block_table, 7)
129
+ assert_block_id(removed, [original_block_ids[1]])
130
+ assert_block_id(block_table, [null_block_id] * 2 + original_block_ids[2:])
131
+
132
+ # 11 tokens are computed. Token 0-7 are out of the sliding window.
133
+ # Block 1002 & 1003 can be removed now. Block 1003 represents a longer
134
+ # sequence, and is expected to be evicted earlier than 1002, so the order
135
+ # of removed blocks should be [1003, 1002].
136
+ removed = manager.remove_skipped_blocks(block_table, 11)
137
+ assert_block_id(removed, [original_block_ids[3], original_block_ids[2]])
138
+ assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
tests/v1/e2e/test_correctness_sliding_window.py
new
+84
@@ -0,0 +1,84 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from dataclasses import dataclass
3
+
4
+import pytest
5
+
6
+from vllm import LLM, SamplingParams
7
+
8
+from ...core.block.e2e.test_correctness_sliding_window import (check_answers,
9
+ prep_prompts)
10
+
11
+
12
+@dataclass
13
+class TestConfig:
14
+ sliding_window: int
15
+ ln_range: tuple[int, int]
16
+
17
+
18
+model_config = {
19
+ "bigcode/starcoder2-3b": TestConfig(4096, (800, 1100)),
20
+ "google/gemma-2-2b-it": TestConfig(4096, (400, 800)),
21
+}
22
+
23
+
24
+@pytest.mark.parametrize(
25
+ "model",
26
+ [
27
+ "bigcode/starcoder2-3b", # sliding window only
28
+ "google/gemma-2-2b-it", # sliding window + full attention
29
+ ])
30
+@pytest.mark.parametrize("batch_size", [5])
31
+@pytest.mark.parametrize("seed", [1])
32
+def test_sliding_window_retrival(monkeypatch, model, batch_size, seed):
33
+ """
34
+ The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
35
+ asks for value of one of them (which is outside the sliding window).
36
+ If we tell it upfront which we are going to be looking for, then
37
+ it answers correctly (mostly).
38
+ """
39
+ with monkeypatch.context() as m:
40
+ m.setenv("VLLM_USE_V1", "1")
41
+
42
+ test_config = model_config[model]
43
+
44
+ llm = LLM(model=model)
45
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
46
+
47
+ prompts, answer, indices = prep_prompts(batch_size,
48
+ ln_range=test_config.ln_range)
49
+
50
+ check_length(prompts, llm, test_config.sliding_window)
51
+
52
+ # Fresh generation
53
+ responses = llm.generate(prompts, sampling_params)
54
+ check_answers(indices,
55
+ answer,
56
+ [response.outputs[0].text for response in responses],
57
+ accept_rate=1.0)
58
+
59
+ # Re-generate with the same prompts to test prefix caching
60
+ responses = llm.generate(prompts, sampling_params)
61
+ check_answers(indices,
62
+ answer,
63
+ [response.outputs[0].text for response in responses],
64
+ accept_rate=1.0)
65
+
66
+
67
+def check_length(prompts: list[str], llm: LLM, sliding_window: int):
68
+ """
69
+ Check if the prompt length is valid, i.e., longer than the sliding window
70
+ size and shorter than the model's max length.
71
+
72
+ Args:
73
+ prompts: list of prompts
74
+ llm: LLM object
75
+ sliding_window: Sliding window size
76
+ """
77
+ tokenizer = llm.get_tokenizer()
78
+ max_model_len = llm.llm_engine.model_config.max_model_len
79
+ assert any(
80
+ len(tokenizer.encode(prompt)) > sliding_window
81
+ for prompt in prompts), "Prompt is too short for test"
82
+ assert all(
83
+ len(tokenizer.encode(prompt)) <= max_model_len
84
+ for prompt in prompts), "Prompt is too long for test"
vllm/config.py
+1
-2
@@ -1116,8 +1116,7 @@ class CacheConfig:
1116
is_attention_free: Whether the model is attention-free.
1117
num_gpu_blocks_override: Number of GPU blocks to use. This overrides the
1118
profiled num_gpu_blocks if specified. Does nothing if None.
1119
- sliding_window: Sliding window size for the KV cache. Can not work with
1120
- prefix caching enabled.
1119
+ sliding_window: Sliding window size for the KV cache.
1120
enable_prefix_caching: Whether to enable prefix caching.
1121
cpu_offload_gb: Size of the CPU offload buffer in GiB.
1122
"""
vllm/v1/core/block_pool.py
+11
-4
@@ -27,6 +27,7 @@ class BlockPool:
27
"""
28
29
def __init__(self, num_gpu_blocks: int, enable_caching: bool):
30
+ assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
31
self.num_gpu_blocks = num_gpu_blocks
32
self.enable_caching = enable_caching
33
# All kv-cache blocks.
@@ -50,6 +51,11 @@ class BlockPool:
51
self.cached_block_hash_to_block: dict[BlockHashType, dict[
52
int, KVCacheBlock]] = defaultdict(dict)
53
54
+ # To represent a placeholder block with block_id=0.
55
+ # The ref_cnt of null_block is not maintained, needs special care to
56
+ # avoid freeing it.
57
+ self.null_block = self.free_block_queue.popleft()
58
+
59
def get_cached_block(self,
60
block_hash: BlockHashType) -> Optional[KVCacheBlock]:
61
"""Get a cached block by the block hash, or None if cache miss.
@@ -214,7 +220,7 @@ class BlockPool:
220
for block in blocks:
221
# ref_cnt=0 means this block is in the free list (i.e. eviction
222
# candidate), so remove it.
217
- if block.ref_cnt == 0:
223
+ if block.ref_cnt == 0 and block != self.null_block:
224
self.free_block_queue.remove(block)
225
block.incr_ref()
226
@@ -228,7 +234,8 @@ class BlockPool:
234
"""
235
for block in ordered_blocks:
236
block.decr_ref()
231
- if block.ref_cnt == 0:
237
+ # null_block should not be added to the free list.
238
+ if block.ref_cnt == 0 and block != self.null_block:
239
self.free_block_queue.append(block)
240
241
def reset_prefix_cache(self) -> bool:
@@ -241,10 +248,10 @@ class BlockPool:
248
False otherwise.
249
"""
250
num_used_blocks = (self.num_gpu_blocks - self.get_num_free_blocks())
244
- if num_used_blocks > 0:
251
+ if num_used_blocks != 1: # The null block is always marked as used
252
logger.warning(
253
"Failed to reset prefix cache because some "
247
- "blocks (%d) are not freed yet", num_used_blocks)
254
+ "blocks (%d) are not freed yet", num_used_blocks - 1)
255
return False
256
257
# Remove all hashes so that no new blocks will hit.
vllm/v1/core/kv_cache_manager.py
+50
-21
@@ -9,6 +9,8 @@ from vllm.utils import cdiv, sha256
9
from vllm.v1.core.block_pool import BlockPool
10
from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
11
hash_request_tokens)
12
+from vllm.v1.core.specialized_manager import get_specialized_manager
13
+from vllm.v1.kv_cache_interface import KVCacheConfig
14
from vllm.v1.metrics.stats import PrefixCacheStats
15
from vllm.v1.request import Request, RequestStatus
16
@@ -19,20 +21,22 @@ class KVCacheManager:
21
22
def __init__(
23
self,
22
- block_size: int,
23
- num_gpu_blocks: int,
24
+ kv_cache_config: KVCacheConfig,
25
max_model_len: int,
25
- sliding_window: Optional[int] = None,
26
enable_caching: bool = True,
27
caching_hash_algo: str = "builtin",
28
num_preallocate_tokens: int = 64,
29
log_stats: bool = False,
30
) -> None:
31
- self.block_size = block_size
32
- self.num_gpu_blocks = num_gpu_blocks
31
+ assert len(kv_cache_config.kv_cache_groups) == 1, (
32
+ "KVCacheManager does not support hybrid models with more than 1 "
33
+ "kv cache group")
34
+ kv_cache_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
35
+ self.block_size = kv_cache_spec.block_size
36
+ self.num_gpu_blocks = kv_cache_config.num_blocks
37
self.max_model_len = max_model_len
34
- self.max_num_blocks_per_req = cdiv(max_model_len, block_size)
35
- self.sliding_window = sliding_window
38
+ self.max_num_blocks_per_req = cdiv(max_model_len, self.block_size)
39
+
40
self.enable_caching = enable_caching
41
self.caching_hash_fn = sha256 if caching_hash_algo == "sha256" else hash
42
# FIXME: make prefix cache stats conditional on log_stats
@@ -48,9 +52,15 @@ class KVCacheManager:
52
# further allocation. When it uses up all the N empty blocks, it gets
53
# N new empty blocks.
54
self.num_preallocate_tokens = num_preallocate_tokens
51
- self.num_preallocate_blocks = cdiv(num_preallocate_tokens, block_size)
55
+ self.num_preallocate_blocks = cdiv(num_preallocate_tokens,
56
+ self.block_size)
57
+
58
+ self.block_pool = BlockPool(self.num_gpu_blocks, enable_caching)
59
53
- self.block_pool = BlockPool(num_gpu_blocks, enable_caching)
60
+ self.specialized_manager = get_specialized_manager(
61
+ kv_cache_spec=kv_cache_spec,
62
+ block_pool=self.block_pool,
63
+ )
64
65
# Mapping from request ID to blocks to track the blocks allocated
66
# for each request, so that we can free the blocks when the request
@@ -117,17 +127,25 @@ class KVCacheManager:
127
128
self.prefix_cache_stats.requests += 1
129
if request.sampling_params.prompt_logprobs is None:
120
- # Check for cache hits
121
- computed_blocks = []
122
- for block_hash in block_hashes:
123
- # block_hashes is a chain of block hashes. If a block hash
124
- # is not in the cached_block_hash_to_id, the following
125
- # block hashes are not computed yet for sure.
126
- if cached_block := self.block_pool.get_cached_block(
127
- block_hash):
128
- computed_blocks.append(cached_block)
129
- else:
130
- break
130
+ if len(block_hashes) * self.block_size == request.num_tokens:
131
+ # When prompt length is divisible by the block size and all
132
+ # blocks are cached, we need to recompute the last token. This
133
+ # have to be achieved by re-computing an entire block because
134
+ # allocate_slots() assumes num_computed_tokens is always a
135
+ # multiple of the block size. To achieve this, remove the last
136
+ # block hash from the block_hashes for find_longest_cache_hit
137
+ # This limitation can potentially be removed in the future to
138
+ # slightly improve the performance.
139
+ last_block_hash = block_hashes.pop()
140
+ else:
141
+ last_block_hash = None
142
+
143
+ computed_blocks = (
144
+ self.specialized_manager.find_longest_cache_hit(block_hashes))
145
+
146
+ if last_block_hash is not None:
147
+ # Add back the last block hash if it was removed.
148
+ block_hashes.append(last_block_hash)
149
150
self.prefix_cache_stats.queries += len(block_hashes)
151
self.prefix_cache_stats.hits += len(computed_blocks)
@@ -176,13 +194,24 @@ class KVCacheManager:
194
195
new_computed_blocks = new_computed_blocks or []
196
197
+ req_blocks = self.req_to_blocks[request.request_id]
198
+
199
+ # Free the blocks that are skipped during the attention computation
200
+ # (e.g., tokens outside the sliding window).
201
+ # We can do this even if we cannot schedule this request due to
202
+ # insufficient free blocks.
203
+ # Should call this function before allocating new blocks to reduce
204
+ # the number of evicted blocks.
205
+ removed_blocks = self.specialized_manager.remove_skipped_blocks(
206
+ req_blocks, request.num_computed_tokens)
207
+ self.block_pool.free_blocks(removed_blocks)
208
+
209
# The number of computed tokens is the number of computed tokens plus
210
# the new prefix caching hits
211
num_computed_tokens = (request.num_computed_tokens +
212
len(new_computed_blocks) * self.block_size)
213
num_required_blocks = cdiv(num_computed_tokens + num_tokens,
214
self.block_size)
185
- req_blocks = self.req_to_blocks[request.request_id]
215
num_new_blocks = (num_required_blocks - len(req_blocks) -
216
len(new_computed_blocks))
217
vllm/v1/core/kv_cache_utils.py
+32
-3
@@ -9,8 +9,9 @@ from typing import Any, Callable, NamedTuple, Optional
9
from vllm.config import VllmConfig
10
from vllm.logger import init_logger
11
from vllm.utils import sha256
12
-from vllm.v1.kv_cache_interface import (KVCacheConfig, KVCacheGroupSpec,
13
- KVCacheSpec, KVCacheTensor)
12
+from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
13
+ KVCacheGroupSpec, KVCacheSpec,
14
+ KVCacheTensor, SlidingWindowSpec)
15
from vllm.v1.metrics.stats import PrefixCacheStats
16
from vllm.v1.request import Request
17
@@ -483,7 +484,7 @@ def check_enough_kv_cache_memory(vllm_config: VllmConfig,
484
max_model_len = vllm_config.model_config.max_model_len
485
needed_memory = 0
486
for layer_spec in kv_cache_spec.values():
486
- needed_memory += layer_spec.bytes_for_tokens(max_model_len)
487
+ needed_memory += layer_spec.max_memory_usage_bytes(vllm_config)
488
489
if needed_memory > available_memory:
490
raise ValueError(
@@ -597,6 +598,33 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
598
return kv_cache_config
599
600
601
+def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
602
+ """
603
+ Only models with one type of KV cache are supported yet. This function tries
604
+ to convert the KV cache specs to one type if the model is a hybrid model
605
+ with multiple type of KV cache. It will convert all SlidingWindowSpec to
606
+ FullAttentionSpec if both types are present.
607
+
608
+ Args:
609
+ kv_cache_spec: The kv cache spec of each attention layer in the model
610
+ """
611
+
612
+ has_full_attention = any(
613
+ isinstance(spec, FullAttentionSpec) for spec in kv_cache_spec.values())
614
+ has_sliding_window = any(
615
+ isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values())
616
+ if has_full_attention and has_sliding_window:
617
+ for layer_name, spec in kv_cache_spec.items():
618
+ if isinstance(spec, SlidingWindowSpec):
619
+ kv_cache_spec[layer_name] = FullAttentionSpec(
620
+ block_size=spec.block_size,
621
+ num_kv_heads=spec.num_kv_heads,
622
+ head_size=spec.head_size,
623
+ dtype=spec.dtype,
624
+ use_mla=spec.use_mla,
625
+ )
626
+
627
+
628
def get_kv_cache_config(vllm_config: VllmConfig,
629
kv_cache_spec: dict[str, KVCacheSpec],
630
available_memory: int) -> KVCacheConfig:
@@ -613,6 +641,7 @@ def get_kv_cache_config(vllm_config: VllmConfig,
641
The generated KVCacheConfigs
642
"""
643
check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
644
+ unify_hybrid_kv_cache_specs(kv_cache_spec)
645
if is_kv_cache_type_uniform(kv_cache_spec):
646
# KV cache of all layers are the same, which is true for
647
# most models. Allocate the same amount of memory for
vllm/v1/core/sched/scheduler.py
+5
-17
@@ -19,6 +19,7 @@ from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
19
from vllm.v1.core.sched.utils import check_stop
20
from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,
21
EngineCoreOutputs)
22
+from vllm.v1.kv_cache_interface import KVCacheConfig
23
from vllm.v1.metrics.stats import SchedulerStats
24
from vllm.v1.outputs import ModelRunnerOutput
25
from vllm.v1.request import Request, RequestStatus
@@ -35,6 +36,7 @@ class Scheduler(SchedulerInterface):
36
model_config: ModelConfig,
37
cache_config: CacheConfig,
38
lora_config: Optional[LoRAConfig],
39
+ kv_cache_config: KVCacheConfig,
40
structured_output_manager: StructuredOutputManager,
41
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
42
include_finished_set: bool = False,
@@ -43,6 +45,7 @@ class Scheduler(SchedulerInterface):
45
self.scheduler_config = scheduler_config
46
self.cache_config = cache_config
47
self.lora_config = lora_config
48
+ self.kv_cache_config = kv_cache_config
49
self.log_stats = log_stats
50
self.structured_output_manager = structured_output_manager
51
@@ -58,15 +61,11 @@ class Scheduler(SchedulerInterface):
61
self.scheduler_config.max_num_batched_tokens
62
self.max_model_len = self.scheduler_config.max_model_len
63
61
- num_gpu_blocks = cache_config.num_gpu_blocks
62
- assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
64
# Create the KV cache manager.
65
self.kv_cache_manager = KVCacheManager(
65
- block_size=self.cache_config.block_size,
66
- num_gpu_blocks=num_gpu_blocks,
66
+ kv_cache_config=kv_cache_config,
67
max_model_len=self.max_model_len,
68
- sliding_window=self.cache_config.sliding_window,
69
- enable_caching=self.cache_config.enable_prefix_caching,
68
+ enable_caching=cache_config.enable_prefix_caching,
69
caching_hash_algo=self.cache_config.prefix_caching_hash_algo,
70
log_stats=self.log_stats)
71
self.block_size = self.cache_config.block_size
@@ -300,17 +299,6 @@ class Scheduler(SchedulerInterface):
299
# `request.num_prompt_tokens` to consider the resumed requests,
300
# which have output tokens.
301
num_new_tokens = request.num_tokens - num_computed_tokens
303
- if num_new_tokens == 0:
304
- # This happens when prompt length is divisible by the block
305
- # size and all blocks are cached. Now we force to recompute
306
- # the last block. Note that we have to re-compute an entire
307
- # block because allocate_slots() assumes num_computed_tokens
308
- # is always a multiple of the block size. This limitation
309
- # can potentially be removed in the future to slightly
310
- # improve the performance.
311
- num_computed_tokens -= self.block_size
312
- num_new_tokens = self.block_size
313
- computed_blocks.pop()
302
if (0 < self.scheduler_config.long_prefill_token_threshold <
303
num_new_tokens):
304
num_new_tokens = (
vllm/v1/core/specialized_manager.py
new
+161
@@ -0,0 +1,161 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from abc import ABC, abstractmethod
3
+
4
+from vllm.utils import cdiv
5
+from vllm.v1.core.block_pool import BlockPool
6
+from vllm.v1.core.kv_cache_utils import BlockHashType, KVCacheBlock
7
+from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheSpec,
8
+ SlidingWindowSpec)
9
+
10
+
11
+class SpecializedManager(ABC):
12
+ """
13
+ An abstract base class for specialized managers that handle the kv
14
+ cache management logic of different attention layers.
15
+ """
16
+
17
+ def __init__(
18
+ self,
19
+ kv_cache_spec: KVCacheSpec,
20
+ block_pool: BlockPool,
21
+ ) -> None:
22
+ """
23
+ Initializes the SpecializedManager.
24
+ Args:
25
+ kv_cache_spec: The kv_cache_spec for this manager.
26
+ block_pool: The block pool.
27
+ """
28
+
29
+ self.block_size = kv_cache_spec.block_size
30
+ self.kv_cache_spec = kv_cache_spec
31
+ self.block_pool = block_pool
32
+
33
+ @abstractmethod
34
+ def find_longest_cache_hit(
35
+ self, block_hashes: list[BlockHashType]) -> list[KVCacheBlock]:
36
+ """
37
+ Get the longest cache hit prefix of the blocks. If no cache hit is
38
+ found, return an empty list.
39
+
40
+ Args:
41
+ block_hashes: The block hashes of the request.
42
+ Returns:
43
+ A list of cached blocks with skipped blocks replaced by null block.
44
+ For example, sliding window manager should return a list like
45
+ [NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)] for block size 4 and
46
+ sliding window 8.
47
+ """
48
+
49
+ raise NotImplementedError
50
+
51
+ @abstractmethod
52
+ def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
53
+ num_computed_tokens: int) -> list[KVCacheBlock]:
54
+ """
55
+ Remove the blocks that are no longer needed from `blocks`. The removed
56
+ blocks should be replaced by null_block. Return the removed blocks in
57
+ eviction order, where the first returned block should be evicted first.
58
+ Don't free the removed blocks in this function.
59
+
60
+ Args:
61
+ blocks: The list of blocks to be updated.
62
+ num_computed_tokens: The number of tokens that have been computed.
63
+ Returns:
64
+ The removed blocks in eviction order.
65
+ """
66
+ raise NotImplementedError
67
+
68
+
69
+class FullAttentionManager(SpecializedManager):
70
+
71
+ def find_longest_cache_hit(
72
+ self, block_hashes: list[BlockHashType]) -> list[KVCacheBlock]:
73
+ computed_blocks: list[KVCacheBlock] = []
74
+ for block_hash in block_hashes:
75
+ # block_hashes is a chain of block hashes. If a block hash is not
76
+ # in the cached_block_hash_to_id, the following block hashes are
77
+ # not computed yet for sure.
78
+ if cached_block := self.block_pool.get_cached_block(block_hash):
79
+ computed_blocks.append(cached_block)
80
+ else:
81
+ break
82
+ return computed_blocks
83
+
84
+ def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
85
+ num_computed_tokens: int) -> list[KVCacheBlock]:
86
+ # No need to remove blocks for full attention.
87
+ return []
88
+
89
+
90
+class SlidingWindowManager(SpecializedManager):
91
+
92
+ def __init__(self, kv_cache_spec: SlidingWindowSpec,
93
+ block_pool: BlockPool):
94
+ super().__init__(kv_cache_spec, block_pool)
95
+ self.sliding_window = kv_cache_spec.sliding_window
96
+ # The number of contiguous blocks needed for prefix cache hit.
97
+ # -1 since the input token itself is also included in the window
98
+ self.sliding_window_contiguous_blocks = cdiv(
99
+ (kv_cache_spec.sliding_window - 1), self.block_size)
100
+ self._null_block = block_pool.null_block
101
+
102
+ def find_longest_cache_hit(
103
+ self, block_hashes: list[BlockHashType]) -> list[KVCacheBlock]:
104
+ # TODO: reduce i by sliding_window_contiguous_blocks when cache miss, to
105
+ # optimize the time complexity from O(len(block_hashes)) to
106
+ # O(len(block_hashes) / sliding_window_contiguous_blocks +
107
+ # sliding_window_contiguous_blocks),
108
+ # which is good for low cache hit rate scenarios.
109
+ computed_blocks = [self._null_block] * len(block_hashes)
110
+ num_contiguous_blocks = 0
111
+
112
+ # Search from right to left and early stop when a match is found.
113
+ for i in range(len(block_hashes) - 1, -1, -1):
114
+ if cached_block := self.block_pool.get_cached_block(
115
+ block_hashes[i]):
116
+ computed_blocks[i] = cached_block
117
+ num_contiguous_blocks += 1
118
+ if (num_contiguous_blocks
119
+ >= self.sliding_window_contiguous_blocks):
120
+ # Trim the trailing blocks.
121
+ # E.g., [NULL, NULL, 8, 3, NULL, 9] -> [NULL, NULL, 8, 3]
122
+ # when sliding_window_contiguous_blocks=2.
123
+ del computed_blocks[i + num_contiguous_blocks:]
124
+ return computed_blocks
125
+ else:
126
+ num_contiguous_blocks = 0
127
+ # The first `num_contiguous_blocks` is a cache hit even if
128
+ # `num_contiguous_blocks < sliding_window_contiguous_blocks`.
129
+ del computed_blocks[num_contiguous_blocks:]
130
+ return computed_blocks
131
+
132
+ def remove_skipped_blocks(self, blocks: list[KVCacheBlock],
133
+ num_computed_tokens: int) -> list[KVCacheBlock]:
134
+ # Remove the blocks that are no longer be in the sliding window and
135
+ # skipped during the attention computation.
136
+ last_useful_token = num_computed_tokens - self.sliding_window + 1
137
+ last_useful_block = last_useful_token // self.block_size
138
+
139
+ removed_blocks: list[KVCacheBlock] = []
140
+ for i in range(last_useful_block - 1, -1, -1):
141
+ if blocks[i] == self._null_block:
142
+ # If the block is already a null block, the blocks before it
143
+ # should also have been set to null blocks by the previous calls
144
+ # to this function.
145
+ break
146
+ removed_blocks.append(blocks[i])
147
+ blocks[i] = self._null_block
148
+ return removed_blocks
149
+
150
+
151
+spec_manager_map: dict[type[KVCacheSpec], type[SpecializedManager]] = {
152
+ FullAttentionSpec: FullAttentionManager,
153
+ SlidingWindowSpec: SlidingWindowManager,
154
+}
155
+
156
+
157
+def get_specialized_manager(kv_cache_spec: KVCacheSpec,
158
+ block_pool: BlockPool) -> SpecializedManager:
159
+ manager_class = spec_manager_map[type(kv_cache_spec)]
160
+ manager = manager_class(kv_cache_spec, block_pool)
161
+ return manager
vllm/v1/engine/core.py
+11
-7
@@ -33,6 +33,7 @@ from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
33
EngineCoreRequestType, UtilityOutput)
34
from vllm.v1.engine.mm_input_cache import MMInputCacheServer
35
from vllm.v1.executor.abstract import Executor
36
+from vllm.v1.kv_cache_interface import KVCacheConfig
37
from vllm.v1.outputs import ModelRunnerOutput
38
from vllm.v1.request import Request, RequestStatus
39
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
@@ -66,8 +67,9 @@ class EngineCore:
67
self.model_executor = executor_class(vllm_config)
68
69
# Setup KV Caches and update CacheConfig after profiling.
69
- num_gpu_blocks, num_cpu_blocks = self._initialize_kv_caches(
70
- vllm_config)
70
+ num_gpu_blocks, num_cpu_blocks, kv_cache_config = \
71
+ self._initialize_kv_caches(vllm_config)
72
+
73
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
74
vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks
75
@@ -95,10 +97,11 @@ class EngineCore:
97
model_config=vllm_config.model_config,
98
cache_config=vllm_config.cache_config,
99
lora_config=vllm_config.lora_config,
100
+ kv_cache_config=kv_cache_config,
101
+ structured_output_manager=self.structured_output_manager,
102
include_finished_set=vllm_config.parallel_config.data_parallel_size
103
> 1,
104
log_stats=self.log_stats,
101
- structured_output_manager=self.structured_output_manager,
105
)
106
107
# Setup MM Input Mapper.
@@ -117,8 +120,8 @@ class EngineCore:
120
self.batch_queue_size)
121
self.batch_queue = queue.Queue(self.batch_queue_size)
122
120
- def _initialize_kv_caches(self,
121
- vllm_config: VllmConfig) -> tuple[int, int]:
123
+ def _initialize_kv_caches(
124
+ self, vllm_config: VllmConfig) -> tuple[int, int, KVCacheConfig]:
125
start = time.time()
126
127
# Get all kv cache needed by the model
@@ -143,13 +146,14 @@ class EngineCore:
146
unify_kv_cache_configs(kv_cache_configs)
147
148
# All workers have the same kv_cache_config except layer names, so use
146
- # an arbitrary one to get the number of blocks.
149
+ # an arbitrary one to initialize the scheduler.
150
assert all([
151
cfg.num_blocks == kv_cache_configs[0].num_blocks
152
for cfg in kv_cache_configs
153
])
154
num_gpu_blocks = kv_cache_configs[0].num_blocks
155
num_cpu_blocks = 0
156
+ scheduler_kv_cache_config = kv_cache_configs[0]
157
158
# Initialize kv cache and warmup the execution
159
self.model_executor.initialize_from_config(kv_cache_configs)
@@ -157,7 +161,7 @@ class EngineCore:
161
elapsed = time.time() - start
162
logger.info(("init engine (profile, create kv cache, "
163
"warmup model) took %.2f seconds"), elapsed)
160
- return num_gpu_blocks, num_cpu_blocks
164
+ return num_gpu_blocks, num_cpu_blocks, scheduler_kv_cache_config
165
166
def add_request(self, request: EngineCoreRequest):
167
"""Add request to the scheduler."""
vllm/v1/kv_cache_interface.py
+46
-11
@@ -4,6 +4,7 @@ from dataclasses import dataclass
4
5
import torch
6
7
+from vllm.config import VllmConfig
8
from vllm.logger import init_logger
9
from vllm.utils import cdiv, get_dtype_size
10
@@ -43,28 +44,23 @@ class KVCacheSpec:
44
"""
45
raise NotImplementedError
46
46
- def bytes_for_tokens(self, num_tokens: int) -> int:
47
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
48
"""
48
- The KV cache size for `num_tokens` tokens in bytes. Returns the real
49
- memory size after padding `num_tokens` to full blocks.
49
+ The maximum possible memory usage of this KV cache in bytes.
50
51
Returns:
52
- The KV cache size
52
+ The KV cache size in bytes
53
"""
54
raise NotImplementedError
55
56
57
@dataclass
58
-class FullAttentionSpec(KVCacheSpec):
58
+class AttentionSpec(KVCacheSpec):
59
num_kv_heads: int
60
head_size: int
61
dtype: torch.dtype
62
use_mla: bool
63
64
- @property
65
- def type_id(self) -> str:
66
- return f"full_attention_{self.block_size}_{self.page_size_bytes}"
67
-
64
@property
65
def page_size_bytes(self) -> int:
66
# For MLA we only store a single latent vector
@@ -72,8 +68,47 @@ class FullAttentionSpec(KVCacheSpec):
68
return coef * self.block_size * self.num_kv_heads * self.head_size \
69
* get_dtype_size(self.dtype)
70
75
- def bytes_for_tokens(self, num_tokens: int) -> int:
76
- return cdiv(num_tokens, self.block_size) * self.page_size_bytes
71
+
72
+@dataclass
73
+class FullAttentionSpec(AttentionSpec):
74
+
75
+ @property
76
+ def type_id(self) -> str:
77
+ return f"full_attention_{self.block_size}_{self.page_size_bytes}"
78
+
79
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
80
+ max_model_len = vllm_config.model_config.max_model_len
81
+ return cdiv(max_model_len, self.block_size) * self.page_size_bytes
82
+
83
+
84
+@dataclass
85
+class SlidingWindowSpec(AttentionSpec):
86
+ sliding_window: int
87
+
88
+ def __post_init__(self):
89
+ assert not self.use_mla, "MLA is not supported for sliding window"
90
+
91
+ @property
92
+ def type_id(self) -> str:
93
+ return f"sliding_window_{self.sliding_window}_{self.block_size}_{self.page_size_bytes}" # noqa
94
+
95
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
96
+ max_model_len = vllm_config.model_config.max_model_len
97
+ max_num_batched_tokens = (
98
+ vllm_config.scheduler_config.max_num_batched_tokens)
99
+
100
+ # During chunked prefill, we allocate KV cache for the last
101
+ # `self.sliding_window-1` computed tokens plus the newly scheduled
102
+ # tokens. And we won't allocate KV cache for more than `max_model_len`
103
+ # tokens.
104
+ num_tokens = min(self.sliding_window - 1 + max_num_batched_tokens,
105
+ max_model_len)
106
+
107
+ # +1 here because the sliding window may not start from the beginning
108
+ # of the block. For example, if the block size is 4 and num_token
109
+ # is 4, we need two blocks [XXCD] [EF] to store the sliding
110
+ # window [CDEF] of 6 tokens.
111
+ return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
112
113
114
@dataclass
vllm/v1/worker/gpu_model_runner.py
+19
-9
@@ -28,8 +28,9 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
28
check_use_alibi, is_pin_memory_available)
29
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
30
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
31
-from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
32
- KVCacheSpec)
31
+from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
32
+ KVCacheConfig, KVCacheSpec,
33
+ SlidingWindowSpec)
34
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
35
ModelRunnerOutput)
36
from vllm.v1.sample.metadata import SamplingMetadata
@@ -1572,7 +1573,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1573
# different GPUs, and `kv_cache_config.num_blocks` is set to
1574
# the min of all `num_blocks`. Verify it here.
1575
assert num_blocks >= kv_cache_config.num_blocks
1575
- if isinstance(kv_cache_spec, FullAttentionSpec):
1576
+ if isinstance(kv_cache_spec, AttentionSpec):
1577
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
1578
num_blocks, kv_cache_spec.block_size,
1579
kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
@@ -1611,12 +1612,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1612
# cross-attention
1613
assert isinstance(attn_module, Attention)
1614
if attn_module.attn_type == AttentionType.DECODER:
1614
- kv_cache_spec[layer_name] = FullAttentionSpec(
1615
- block_size=block_size,
1616
- num_kv_heads=attn_module.num_kv_heads,
1617
- head_size=attn_module.head_size,
1618
- dtype=self.kv_cache_dtype,
1619
- use_mla=use_mla)
1615
+ if attn_module.sliding_window is not None:
1616
+ kv_cache_spec[layer_name] = SlidingWindowSpec(
1617
+ block_size=block_size,
1618
+ num_kv_heads=attn_module.num_kv_heads,
1619
+ head_size=attn_module.head_size,
1620
+ dtype=self.kv_cache_dtype,
1621
+ sliding_window=attn_module.sliding_window,
1622
+ use_mla=use_mla)
1623
+ else:
1624
+ kv_cache_spec[layer_name] = FullAttentionSpec(
1625
+ block_size=block_size,
1626
+ num_kv_heads=attn_module.num_kv_heads,
1627
+ head_size=attn_module.head_size,
1628
+ dtype=self.kv_cache_dtype,
1629
+ use_mla=use_mla)
1630
elif attn_module.attn_type in (AttentionType.ENCODER,
1631
AttentionType.ENCODER_ONLY):
1632
# encoder-only attention does not need KV cache.
vllm/v1/worker/tpu_model_runner.py
+18
-10
@@ -29,7 +29,7 @@ from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,
29
PallasMetadata)
30
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
31
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
32
- KVCacheSpec)
32
+ KVCacheSpec, SlidingWindowSpec)
33
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
34
ModelRunnerOutput, SamplerOutput)
35
from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
@@ -353,17 +353,25 @@ class TPUModelRunner:
353
block_size = self.vllm_config.cache_config.block_size
354
kv_cache_spec: dict[str, KVCacheSpec] = {}
355
for layer_name, attn_module in forward_ctx.items():
356
- # TODO: Support other attention modules, e.g., sliding window,
357
- # cross-attention, MLA.
356
assert isinstance(attn_module, Attention)
357
if attn_module.attn_type == AttentionType.DECODER:
360
- kv_cache_spec[layer_name] = FullAttentionSpec(
361
- block_size=block_size,
362
- num_kv_heads=attn_module.num_kv_heads,
363
- head_size=attn_module.head_size,
364
- dtype=attn_module.dtype,
365
- use_mla=False,
366
- )
358
+ if attn_module.sliding_window is not None:
359
+ kv_cache_spec[layer_name] = SlidingWindowSpec(
360
+ block_size=block_size,
361
+ num_kv_heads=attn_module.num_kv_heads,
362
+ head_size=attn_module.head_size,
363
+ dtype=attn_module.dtype,
364
+ sliding_window=attn_module.sliding_window,
365
+ use_mla=False,
366
+ )
367
+ else:
368
+ kv_cache_spec[layer_name] = FullAttentionSpec(
369
+ block_size=block_size,
370
+ num_kv_heads=attn_module.num_kv_heads,
371
+ head_size=attn_module.head_size,
372
+ dtype=attn_module.dtype,
373
+ use_mla=False,
374
+ )
375
elif attn_module.attn_type in (AttentionType.ENCODER,
376
AttentionType.ENCODER_ONLY):
377
# encoder-only attention does not need KV cache.