@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c280066f9d
Fork 0

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 14, 05:30 UTC

[v1] Move block_hashes from KVCacheManager to Request.block_hashes (#19728)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Aug 16, 2025 at 02:52 UTC c280066f9dad0288a768a6234bea08171c4b88b9
19 files changed +382 -336
tests/v1/core/test_async_scheduler.py
+15 -7
@@ -7,6 +7,7 @@ import pytest
7 from vllm.v1.core.sched.output import SchedulerOutput
8 from vllm.v1.outputs import ModelRunnerOutput
9 from vllm.v1.request import RequestStatus
10 +from vllm.v1.utils import ConstantList
11
12 from .utils import create_requests, create_scheduler
13
@@ -140,7 +141,8 @@ def test_prefix_caching_for_prefill_dedup():
141 requests = create_requests(num_requests=5,
142 num_tokens=num_prompt_tokens,
143 max_tokens=3,
143 - same_prompt=True)
144 + same_prompt=True,
145 + block_size=BLOCK_SIZE)
146 requests_copy = requests.copy()
147
148 # Two requests with the same prompt.
@@ -188,7 +190,8 @@ def test_prefix_caching_for_multi_turn():
190 block_size=BLOCK_SIZE)
191 requests = create_requests(num_requests=5,
192 num_tokens=num_prompt_tokens,
191 - max_tokens=num_output_tokens)
193 + max_tokens=num_output_tokens,
194 + block_size=BLOCK_SIZE)
195
196 for req in requests:
197 scheduler.add_request(req)
@@ -208,14 +211,19 @@ def test_prefix_caching_for_multi_turn():
211
212 # Create next-turn requests whose prompts are the full output of the
213 # previous turn.
211 - next_turn_requests = create_requests(
212 - num_requests=5,
213 - num_tokens=num_prompt_tokens + num_output_tokens,
214 - max_tokens=num_output_tokens,
215 - )
214 + next_turn_requests = create_requests(num_requests=5,
215 + num_tokens=num_prompt_tokens +
216 + num_output_tokens,
217 + max_tokens=num_output_tokens,
218 + block_size=BLOCK_SIZE)
219 for i, req in enumerate(next_turn_requests):
220 req.prompt_token_ids = (requests[i].prompt_token_ids +
221 list(requests[i].output_token_ids))
222 + req._all_token_ids = req.prompt_token_ids.copy()
223 + req.all_token_ids = ConstantList(req._all_token_ids)
224 + req.block_hashes = []
225 + req.block_hashes = req.get_hash_new_full_blocks()
226 +
227 # Schedule the next-turn requests.
228 for req in next_turn_requests:
229 scheduler.add_request(req)
tests/v1/core/test_kv_cache_utils.py
+27 -23
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import importlib
4 -from typing import Optional
4 +from typing import Callable, Optional
5
6 import pytest
7 import torch
@@ -19,7 +19,7 @@ from vllm.v1.core.kv_cache_utils import (
19 FreeKVCacheBlockQueue, KVCacheBlock, PrefixCachingMetrics,
20 estimate_max_model_len, generate_block_hash_extra_keys,
21 get_kv_cache_config, get_max_concurrency_for_kv_cache_config,
22 - hash_block_tokens, hash_request_tokens, init_none_hash,
22 + get_request_block_hasher, hash_block_tokens, init_none_hash,
23 is_kv_cache_type_uniform, unify_kv_cache_configs)
24 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
25 KVCacheGroupSpec, KVCacheTensor,
@@ -33,6 +33,8 @@ from vllm.v1.request import Request
33 def make_request(
34 request_id: str,
35 prompt_token_ids: list[int],
36 + block_size: int = 3,
37 + hash_fn: Callable = hash,
38 mm_positions: Optional[list[PlaceholderRange]] = None,
39 mm_hashes: Optional[list[str]] = None,
40 cache_salt: Optional[str] = None,
@@ -49,18 +51,17 @@ def make_request(
51 mm_item = MultiModalKwargsItem.from_elems([mm_elem])
52 mm_kwargs = [mm_item] * len(mm_positions)
53
52 - return Request(
53 - request_id=request_id,
54 - prompt_token_ids=prompt_token_ids,
55 - multi_modal_kwargs=mm_kwargs,
56 - multi_modal_hashes=mm_hashes,
57 - multi_modal_placeholders=mm_positions,
58 - sampling_params=SamplingParams(max_tokens=17),
59 - pooling_params=None,
60 - eos_token_id=100,
61 - lora_request=None,
62 - cache_salt=cache_salt,
63 - )
54 + return Request(request_id=request_id,
55 + prompt_token_ids=prompt_token_ids,
56 + multi_modal_kwargs=mm_kwargs,
57 + multi_modal_hashes=mm_hashes,
58 + multi_modal_placeholders=mm_positions,
59 + sampling_params=SamplingParams(max_tokens=17),
60 + pooling_params=None,
61 + eos_token_id=100,
62 + lora_request=None,
63 + cache_salt=cache_salt,
64 + block_hasher=get_request_block_hasher(block_size, hash_fn))
65
66
67 def new_kv_cache_spec(block_size=16,
@@ -428,12 +429,14 @@ def test_hash_block_tokens(hash_fn):
429
430
431 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
431 -def test_hash_request_tokens(hash_fn):
432 +def test_request_block_hasher(hash_fn):
433 import vllm.v1.core.kv_cache_utils
434 init_none_hash(hash_fn)
435 request = make_request(
436 request_id="0",
437 prompt_token_ids=[_ for _ in range(6)],
438 + block_size=3,
439 + hash_fn=hash_fn,
440 mm_positions=[
441 PlaceholderRange(offset=0, length=3),
442 PlaceholderRange(offset=3, length=3),
@@ -441,9 +444,7 @@ def test_hash_request_tokens(hash_fn):
444 mm_hashes=["hash1", "hash2"],
445 )
446
444 - block_size = 3
445 - block_hashes = hash_request_tokens(hash_fn, block_size, request)
446 -
447 + block_hashes = request.block_hashes
448 assert len(block_hashes) == 2
449 assert isinstance(block_hashes[0], vllm.v1.core.kv_cache_utils.BlockHash)
450 assert isinstance(block_hashes[1], vllm.v1.core.kv_cache_utils.BlockHash)
@@ -464,6 +465,8 @@ def test_hash_tokens_different_mm_input(hash_fn):
465 request1 = make_request(
466 request_id="0",
467 prompt_token_ids=[_ for _ in range(6)],
468 + block_size=3,
469 + hash_fn=hash_fn,
470 mm_positions=[
471 PlaceholderRange(offset=0, length=3),
472 PlaceholderRange(offset=3, length=3),
@@ -479,9 +482,8 @@ def test_hash_tokens_different_mm_input(hash_fn):
482 ],
483 mm_hashes=["hash3", "hash2"],
484 )
482 - block_size = 3
483 - block_hashes1 = hash_request_tokens(hash_fn, block_size, request1)
484 - block_hashes2 = hash_request_tokens(hash_fn, block_size, request2)
485 + block_hashes1 = request1.block_hashes
486 + block_hashes2 = request2.block_hashes
487 assert block_hashes1[0] != block_hashes2[0]
488 assert block_hashes1[1] != block_hashes2[1]
489
@@ -493,12 +495,13 @@ def test_hash_request_tokens_no_mm_inputs(hash_fn):
495 request = make_request(
496 request_id="0",
497 prompt_token_ids=[_ for _ in range(6)],
498 + block_size=3,
499 + hash_fn=hash_fn,
500 mm_positions=None,
501 mm_hashes=None,
502 )
503
500 - block_size = 3
501 - block_hashes = hash_request_tokens(hash_fn, block_size, request)
504 + block_hashes = request.block_hashes
505
506 assert len(block_hashes) == 2
507 assert block_hashes[0].token_ids == (0, 1, 2)
@@ -858,6 +861,7 @@ def test_allocate_with_lookahead():
861 request = make_request(
862 request_id="0",
863 prompt_token_ids=[],
864 + block_size=block_size,
865 mm_positions=None,
866 mm_hashes=None,
867 )
tests/v1/core/test_prefix_caching.py
+122 -103
@@ -3,7 +3,7 @@
3 """Compare the with and without prefix caching."""
4
5 import copy
6 -from typing import Optional
6 +from typing import Callable, Optional
7
8 import pytest
9 import torch
@@ -17,8 +17,9 @@ from vllm.utils import sha256, sha256_cbor_64bit
17 from vllm.v1.core.block_pool import BlockPool
18 from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
19 from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
20 - KVCacheBlock, hash_block_tokens,
21 - init_none_hash)
20 + KVCacheBlock,
21 + get_request_block_hasher,
22 + hash_block_tokens, init_none_hash)
23 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
24 KVCacheGroupSpec, SlidingWindowSpec)
25
@@ -26,6 +27,8 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
27 def make_request(
28 request_id: str,
29 prompt_token_ids: list[int],
30 + block_size: int,
31 + hash_fn: Callable,
32 mm_positions: Optional[list[PlaceholderRange]] = None,
33 mm_hashes: Optional[list[str]] = None,
34 prompt_logprobs: Optional[int] = None,
@@ -43,19 +46,18 @@ def make_request(
46 mm_item = MultiModalKwargsItem.from_elems([mm_elem])
47 mm_kwargs = [mm_item] * len(mm_positions)
48
46 - return Request(
47 - request_id=request_id,
48 - prompt_token_ids=prompt_token_ids,
49 - multi_modal_kwargs=mm_kwargs,
50 - multi_modal_hashes=mm_hashes,
51 - multi_modal_placeholders=mm_positions,
52 - sampling_params=SamplingParams(max_tokens=17,
53 - prompt_logprobs=prompt_logprobs),
54 - pooling_params=None,
55 - eos_token_id=100,
56 - lora_request=None,
57 - cache_salt=cache_salt,
58 - )
49 + return Request(request_id=request_id,
50 + prompt_token_ids=prompt_token_ids,
51 + multi_modal_kwargs=mm_kwargs,
52 + multi_modal_hashes=mm_hashes,
53 + multi_modal_placeholders=mm_positions,
54 + sampling_params=SamplingParams(
55 + max_tokens=17, prompt_logprobs=prompt_logprobs),
56 + pooling_params=None,
57 + eos_token_id=100,
58 + lora_request=None,
59 + cache_salt=cache_salt,
60 + block_hasher=get_request_block_hasher(block_size, hash_fn))
61
62
63 def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
@@ -105,11 +107,11 @@ def make_kv_cache_config_hybrid_model(block_size: int,
107
108 @pytest.mark.parametrize("hash_algo", ["sha256", "sha256_cbor_64bit", "hash"])
109 def test_prefill(hash_algo):
110 + block_size = 16
111 manager = KVCacheManager(
109 - make_kv_cache_config(16, 11),
112 + make_kv_cache_config(block_size, 11),
113 max_model_len=8192,
114 enable_caching=True,
112 - caching_hash_algo=hash_algo,
115 )
116
117 # choose the hash function according to the parameter
@@ -123,9 +125,9 @@ def test_prefill(hash_algo):
125 # Incomplete 1 block (7 tokens)
126 unique_token_ids = [3] * 7
127 all_token_ids = common_token_ids + unique_token_ids
126 - req0 = make_request("0", all_token_ids)
128 + req0 = make_request("0", all_token_ids, block_size, hash_fn)
129 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
128 - assert len(manager.req_to_block_hashes[req0.request_id]) == 3
130 + assert len(req0.block_hashes) == 3
131 assert not computed_blocks.blocks[0]
132 assert num_computed_tokens == 0
133 blocks = manager.allocate_slots(req0, 55,
@@ -152,9 +154,10 @@ def test_prefill(hash_algo):
154 # Cache hit in the common prefix when the original block is still in use.
155 # Incomplete 1 block (5 tokens)
156 unique_token_ids = [3] * 5
155 - req1 = make_request("1", common_token_ids + unique_token_ids)
157 + req1 = make_request("1", common_token_ids + unique_token_ids, block_size,
158 + hash_fn)
159 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
157 - assert len(manager.req_to_block_hashes[req1.request_id]) == 3
160 + assert len(req1.block_hashes) == 3
161 assert computed_blocks.get_block_ids() == ([1, 2, 3], )
162 assert num_computed_tokens == 3 * 16
163 num_new_tokens = 53 - 3 * 16
@@ -187,9 +190,10 @@ def test_prefill(hash_algo):
190 # Cache hit in the common prefix when the original block is already free.
191 # Incomplete 1 block (6 tokens)
192 unique_token_ids = [3] * 6
190 - req2 = make_request("2", common_token_ids + unique_token_ids)
193 + req2 = make_request("2", common_token_ids + unique_token_ids, block_size,
194 + hash_fn)
195 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
192 - assert len(manager.req_to_block_hashes[req2.request_id]) == 3
196 + assert len(req2.block_hashes) == 3
197 assert computed_blocks.get_block_ids() == ([1, 2, 3], )
198 assert num_computed_tokens == 3 * 16
199 num_new_tokens = 53 - 3 * 16
@@ -208,7 +212,7 @@ def test_prefill(hash_algo):
212 manager.free(req2)
213
214 # Cache miss and eviction.
211 - req3 = make_request("3", [99] * (16 * 10))
215 + req3 = make_request("3", [99] * (16 * 10), block_size, hash_fn)
216 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
217 assert not computed_blocks.blocks[0]
218 assert num_computed_tokens == 0
@@ -242,9 +246,9 @@ def test_prefill_hybrid_model():
246 # Incomplete 1 block (7 tokens)
247 unique_token_ids = [3] * 7
248 all_token_ids = common_token_ids + unique_token_ids
245 - req0 = make_request("0", all_token_ids)
249 + req0 = make_request("0", all_token_ids, block_size, hash_fn)
250 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
247 - assert len(manager.req_to_block_hashes[req0.request_id]) == 3
251 + assert len(req0.block_hashes) == 3
252 assert not computed_blocks.blocks[0]
253 assert num_computed_tokens == 0
254 blocks = manager.allocate_slots(req0, 55,
@@ -274,9 +278,10 @@ def test_prefill_hybrid_model():
278 # Cache hit in the common prefix
279 # Incomplete 1 block (5 tokens)
280 unique_token_ids = [3] * 5
277 - req1 = make_request("1", common_token_ids + unique_token_ids)
281 + req1 = make_request("1", common_token_ids + unique_token_ids, block_size,
282 + hash_fn)
283 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
279 - assert len(manager.req_to_block_hashes[req1.request_id]) == 3
284 + assert len(req1.block_hashes) == 3
285 assert computed_blocks.get_block_ids() == ([1, 2, 3], [0, 6,
286 7], [0, 10, 11])
287 assert num_computed_tokens == 3 * 16
@@ -290,7 +295,7 @@ def test_prefill_hybrid_model():
295 if block != manager.block_pool.null_block:
296 assert block.ref_cnt == 2
297
293 - block_hashes = manager.req_to_block_hashes[req1.request_id]
298 + block_hashes = req1.block_hashes
299 manager.free(req0)
300 manager.free(req1)
301
@@ -300,12 +305,13 @@ def test_prefill_hybrid_model():
305 def test_partial_request_hit(request_id: str,
306 hash_to_evict: list[BlockHashWithGroupId],
307 expect_hit_length: int):
303 - req = make_request(request_id, common_token_ids + unique_token_ids)
308 + req = make_request(request_id, common_token_ids + unique_token_ids,
309 + block_size, hash)
310 for hash_with_group_id in hash_to_evict:
311 manager.block_pool.cached_block_hash_to_block.pop(
312 hash_with_group_id)
313 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
308 - assert len(manager.req_to_block_hashes[req.request_id]) == 3
314 + assert len(req.block_hashes) == 3
315 assert num_computed_tokens == expect_hit_length * block_size
316 for block_per_group in computed_blocks.blocks:
317 assert len(block_per_group) == num_computed_tokens // block_size
@@ -364,8 +370,9 @@ def test_prefill_plp():
370 2. Schedule non-plp request and validate blocks
371 3. Schedule plp request; no hit should occur; validate blocks
372 '''
373 + block_size = 16
374 manager = KVCacheManager(
368 - make_kv_cache_config(16, 11),
375 + make_kv_cache_config(block_size, 11),
376 max_model_len=8192,
377 enable_caching=True,
378 )
@@ -380,9 +387,13 @@ def test_prefill_plp():
387 # Incomplete 1 block (7 tokens)
388 unique_token_ids = [3] * 7
389 all_token_ids = common_token_ids + unique_token_ids
383 - req0 = make_request("0", all_token_ids, prompt_logprobs=5)
390 + req0 = make_request("0",
391 + all_token_ids,
392 + block_size,
393 + hash_fn,
394 + prompt_logprobs=5)
395 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
385 - assert len(manager.req_to_block_hashes[req0.request_id]) == 0
396 + assert len(req0.block_hashes) == 3
397 assert not computed_blocks.blocks[0]
398 assert num_computed_tokens == 0
399 blocks = manager.allocate_slots(req0, 55,
@@ -411,9 +422,10 @@ def test_prefill_plp():
422 # Cache hit in the common prefix when the original block is still in use.
423 # Incomplete 1 block (5 tokens)
424 unique_token_ids = [3] * 5
414 - req1 = make_request("1", common_token_ids + unique_token_ids)
425 + req1 = make_request("1", common_token_ids + unique_token_ids, block_size,
426 + hash_fn)
427 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
416 - assert len(manager.req_to_block_hashes[req1.request_id]) == 3
428 + assert len(req1.block_hashes) == 3
429 assert computed_blocks.get_block_ids() == ([1, 2, 3], )
430 assert num_computed_tokens == 3 * 16
431 num_new_tokens = 53 - 3 * 16
@@ -447,9 +459,11 @@ def test_prefill_plp():
459 unique_token_ids = [3] * 6
460 req2 = make_request("2",
461 common_token_ids + unique_token_ids,
462 + block_size,
463 + hash_fn,
464 prompt_logprobs=5)
465 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
452 - assert len(manager.req_to_block_hashes[req2.request_id]) == 0
466 + assert len(req2.block_hashes) == 3
467 assert not computed_blocks.blocks[0]
468 assert num_computed_tokens == 0
469 blocks = manager.allocate_slots(req2, 55,
@@ -469,8 +483,9 @@ def test_prefill_plp():
483
484
485 def test_decode():
486 + block_size = 16
487 manager = KVCacheManager(
473 - make_kv_cache_config(16, 11),
488 + make_kv_cache_config(block_size, 11),
489 max_model_len=8192,
490 enable_caching=True,
491 )
@@ -481,7 +496,8 @@ def test_decode():
496 # Fully cache miss
497 # Incomplete 1 block (7 tokens)
498 unique_token_ids = [3] * 7
484 - req0 = make_request("0", common_token_ids + unique_token_ids)
499 + req0 = make_request("0", common_token_ids + unique_token_ids, block_size,
500 + hash)
501 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
502 assert not computed_blocks.blocks[0]
503 assert num_computed_tokens == 0
@@ -518,14 +534,15 @@ def test_decode():
534
535
536 def test_evict():
537 + block_size = 16
538 manager = KVCacheManager(
522 - make_kv_cache_config(16, 11),
539 + make_kv_cache_config(block_size, 11),
540 max_model_len=8192,
541 enable_caching=True,
542 )
543
544 last_token_id = 5 * 16 + 7
528 - req0 = make_request("0", list(range(last_token_id)))
545 + req0 = make_request("0", list(range(last_token_id)), block_size, hash)
546 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
547 assert not computed_blocks.blocks[0]
548 assert num_computed_tokens == 0
@@ -536,7 +553,8 @@ def test_evict():
553
554 # 3 blocks.
555 req1 = make_request("1", list(range(last_token_id,
539 - last_token_id + 3 * 16)))
556 + last_token_id + 3 * 16)), block_size,
557 + hash)
558 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
559 assert not computed_blocks.blocks[0]
560 assert num_computed_tokens == 0
@@ -558,7 +576,7 @@ def test_evict():
576 ] == [10, 6, 5, 4, 3, 2, 1, 9, 8, 7]
577
578 # Touch the first 2 blocks.
561 - req2 = make_request("2", list(range(2 * 16 + 3)))
579 + req2 = make_request("2", list(range(2 * 16 + 3)), block_size, hash)
580 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
581 assert computed_blocks.get_block_ids() == ([1, 2], )
582 assert num_computed_tokens == 2 * 16
@@ -583,7 +601,7 @@ def test_hash_block_correct_reuse():
601
602 # Allocate 1 block and cache it.
603 num_tokens = block_size * 1
586 - req = make_request("0", list(range(num_tokens)))
604 + req = make_request("0", list(range(num_tokens)), block_size, hash)
605 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
606 assert not computed_blocks.blocks[0]
607 assert num_computed_tokens == 0
@@ -597,7 +615,7 @@ def test_hash_block_correct_reuse():
615
616 # Allocate a new block that's not full, make sure hash info on the
617 # block is cleared.
600 - req = make_request("1", list(range(num_tokens - 1)))
618 + req = make_request("1", list(range(num_tokens - 1)), block_size, hash)
619 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
620 assert not computed_blocks.blocks[0]
621 assert num_computed_tokens == 0
@@ -624,7 +642,7 @@ def test_computed_blocks_not_evicted():
642
643 # Allocate a block and cache it.
644 num_tokens = block_size * 1
627 - req0 = make_request("0", list(range(num_tokens)))
645 + req0 = make_request("0", list(range(num_tokens)), block_size, hash)
646 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
647 assert not computed_blocks.blocks[0]
648 assert num_computed_tokens == 0
@@ -635,7 +653,8 @@ def test_computed_blocks_not_evicted():
653 assert blocks.blocks[0][0].block_id == 1
654
655 # Allocate another block.
638 - req1 = make_request("1", list(range(num_tokens, num_tokens * 2)))
656 + req1 = make_request("1", list(range(num_tokens, num_tokens * 2)),
657 + block_size, hash)
658 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
659 assert not computed_blocks.blocks[0]
660 assert num_computed_tokens == 0
@@ -651,7 +670,7 @@ def test_computed_blocks_not_evicted():
670
671 # Now if we have a cache hit on the first block, we should evict the second
672 # cached block rather than the first one.
654 - req2 = make_request("2", list(range(num_tokens * 2)))
673 + req2 = make_request("2", list(range(num_tokens * 2)), block_size, hash)
674 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
675 assert len(computed_blocks.blocks[0]) == 1
676 assert computed_blocks.blocks[0][0].block_id == 1
@@ -675,7 +694,8 @@ def test_basic_prefix_caching_disabled():
694 enable_caching=False,
695 )
696
678 - req1 = make_request("1", list(range(10))) # 2 blocks and some more
697 + req1 = make_request("1", list(range(10)), block_size,
698 + hash) # 2 blocks and some more
699
700 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
701 assert not computed_blocks.blocks[0]
@@ -689,7 +709,8 @@ def test_basic_prefix_caching_disabled():
709 manager.free(req1)
710
711 # No caching.
692 - req2 = make_request("2", list(range(16))) # shared prefix
712 + req2 = make_request("2", list(range(16)), block_size,
713 + hash) # shared prefix
714 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
715 assert not computed_blocks.blocks[0]
716 assert num_computed_tokens == 0
@@ -699,7 +720,7 @@ def test_basic_prefix_caching_disabled():
720 assert len(blocks.blocks[0]) == 4
721
722 # New requests should not have any blocks.
702 - req3 = make_request("3", list(range(4)))
723 + req3 = make_request("3", list(range(4)), block_size, hash)
724 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
725 assert not computed_blocks.blocks[0]
726 assert num_computed_tokens == 0
@@ -727,20 +748,17 @@ def test_cache_blocks(hash_fn):
748 # Block 1: [4, 5, 6, 7]
749 # Block 2: [8, 9, 10, 11]
750 # Block 3: [12, 13]
730 - req = make_request("0", list(range(14)))
751 + req = make_request("0", list(range(14)), block_size, hash_fn)
752
753 # Test that blocks are cached correctly for 2 full blocks from the start.
754 blocks = [KVCacheBlock(block_id=i) for i in range(2)]
734 - block_hashes: list[BlockHash] = []
755
756 block_pool.cache_full_blocks(
757 request=req,
758 blocks=blocks,
739 - block_hashes=block_hashes,
759 num_cached_blocks=0,
760 num_full_blocks=2,
761 block_size=block_size,
743 - hash_fn=hash_fn,
762 kv_cache_group_id=0,
763 )
764
@@ -752,11 +770,9 @@ def test_cache_blocks(hash_fn):
770 block_pool.cache_full_blocks(
771 request=req,
772 blocks=blocks,
755 - block_hashes=block_hashes,
773 num_cached_blocks=2,
774 num_full_blocks=3,
775 block_size=block_size,
759 - hash_fn=hash_fn,
776 kv_cache_group_id=0,
777 )
778 assert len(block_pool.cached_block_hash_to_block) == 3
@@ -775,23 +791,20 @@ def test_cache_blocks_multi_group():
791 # Block 1/5: [4, 5, 6, 7]
792 # Block 2/6: [8, 9, 10, 11]
793 # Block 3/7: [12, 13]
778 - req = make_request("0", list(range(14)))
794 + req = make_request("0", list(range(14)), block_size, hash)
795
796 # Cache the blocks for group 0.
797 blocks = [KVCacheBlock(block_id=i) for i in range(2)]
782 - block_hashes: list[BlockHash] = []
798 block_pool.cache_full_blocks(
799 request=req,
800 blocks=blocks,
786 - block_hashes=block_hashes,
801 num_cached_blocks=0,
802 num_full_blocks=2,
803 block_size=block_size,
790 - hash_fn=hash,
804 kv_cache_group_id=0,
805 )
806 assert len(block_pool.cached_block_hash_to_block) == 2
794 - assert len(block_hashes) == 2
807 + assert len(req.block_hashes) == 3
808 assert all([block.block_hash is not None for block in blocks])
809
810 # Cache the blocks for group 1.
@@ -799,38 +812,36 @@ def test_cache_blocks_multi_group():
812 block_pool.cache_full_blocks(
813 request=req,
814 blocks=blocks,
802 - block_hashes=block_hashes,
815 num_cached_blocks=0,
816 num_full_blocks=3,
817 block_size=block_size,
806 - hash_fn=hash,
818 kv_cache_group_id=1,
819 )
820 assert len(block_pool.cached_block_hash_to_block) == 5
810 - assert len(block_hashes) == 3
821 + assert len(req.block_hashes) == 3
822 assert all([block.block_hash is not None for block in blocks])
823
824 # Block hash 0: hit for group 0 and 1
825 # Block hash 1: hit for group 0 and 1
826 # Block hash 2: hit for group 1
827
817 - assert block_pool.get_cached_block(block_hashes[0],
828 + assert block_pool.get_cached_block(req.block_hashes[0],
829 kv_cache_group_ids=[0]) is not None
819 - assert block_pool.get_cached_block(block_hashes[1],
830 + assert block_pool.get_cached_block(req.block_hashes[1],
831 kv_cache_group_ids=[0]) is not None
821 - assert block_pool.get_cached_block(block_hashes[2],
832 + assert block_pool.get_cached_block(req.block_hashes[2],
833 kv_cache_group_ids=[0]) is None
823 - assert block_pool.get_cached_block(block_hashes[0],
834 + assert block_pool.get_cached_block(req.block_hashes[0],
835 kv_cache_group_ids=[1]) is not None
825 - assert block_pool.get_cached_block(block_hashes[1],
836 + assert block_pool.get_cached_block(req.block_hashes[1],
837 kv_cache_group_ids=[1]) is not None
827 - assert block_pool.get_cached_block(block_hashes[2],
838 + assert block_pool.get_cached_block(req.block_hashes[2],
839 kv_cache_group_ids=[1]) is not None
829 - assert block_pool.get_cached_block(block_hashes[0],
840 + assert block_pool.get_cached_block(req.block_hashes[0],
841 kv_cache_group_ids=[0, 1]) is not None
831 - assert block_pool.get_cached_block(block_hashes[1],
842 + assert block_pool.get_cached_block(req.block_hashes[1],
843 kv_cache_group_ids=[0, 1]) is not None
833 - assert block_pool.get_cached_block(block_hashes[2],
844 + assert block_pool.get_cached_block(req.block_hashes[2],
845 kv_cache_group_ids=[0, 1]) is None
846
847
@@ -838,8 +849,9 @@ def test_mm_prefix_caching():
849 """
850 This tests that the multi-modal prefix caching is correct.
851 """
852 + block_size = 16
853 manager = KVCacheManager(
842 - make_kv_cache_config(16, 11),
854 + make_kv_cache_config(block_size, 11),
855 max_model_len=8192,
856 enable_caching=True,
857 )
@@ -865,6 +877,8 @@ def test_mm_prefix_caching():
877 mm_hashes = common_mm_hashes + ["ccc"]
878 req0 = make_request("0",
879 all_token_ids,
880 + block_size,
881 + hash,
882 mm_positions=mm_positions,
883 mm_hashes=mm_hashes)
884 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
@@ -872,7 +886,7 @@ def test_mm_prefix_caching():
886 # Completed block should have hashes with extra keys.
887 assert not computed_blocks.blocks[0]
888 assert num_computed_tokens == 0
875 - block_hashes = manager.req_to_block_hashes[req0.request_id]
889 + block_hashes = req0.block_hashes
890 assert len(block_hashes) == 3
891 assert block_hashes[0].extra_keys == ("aaa", )
892 assert block_hashes[1].extra_keys == ("aaa", "bbb")
@@ -905,6 +919,8 @@ def test_mm_prefix_caching():
919 mm_hashes = common_mm_hashes + ["ccc"]
920 req1 = make_request("1",
921 all_token_ids,
922 + block_size,
923 + hash,
924 mm_positions=mm_positions,
925 mm_hashes=mm_hashes)
926 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
@@ -927,13 +943,13 @@ def test_cache_key_salting():
943 # 3 complete blocks and an incomplete block with 11 tokens.
944 common_token_ids = [i for i in range(3) for _ in range(block_size)]
945 token_ids = common_token_ids + [3] * 11
930 - req0 = make_request("0", token_ids, cache_salt="salt1")
946 + req0 = make_request("0", token_ids, block_size, hash, cache_salt="salt1")
947 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
948
949 # Completed block should have hashes with extra keys.
950 assert not computed_blocks.blocks[0]
951 assert num_computed_tokens == 0
936 - block_hashes = manager.req_to_block_hashes[req0.request_id]
952 + block_hashes = req0.block_hashes
953 assert len(block_hashes) == 3
954 assert block_hashes[0].extra_keys == ("salt1", )
955 assert block_hashes[1].extra_keys is None
@@ -959,7 +975,7 @@ def test_cache_key_salting():
975
976 # Test cache hit with a new request that has the same salt.
977 token_ids = common_token_ids + [4] * 11
962 - req1 = make_request("1", token_ids, cache_salt="salt1")
978 + req1 = make_request("1", token_ids, block_size, hash, cache_salt="salt1")
979 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
980 # Should match only a prefix of 3 blocks.
981 assert len(computed_blocks.blocks[0]) == 3
@@ -967,11 +983,11 @@ def test_cache_key_salting():
983
984 # Test cache miss with same content but different salt.
985 token_ids = common_token_ids + [4] * 11
970 - req2 = make_request("2", token_ids, cache_salt="salt2")
986 + req2 = make_request("2", token_ids, block_size, hash, cache_salt="salt2")
987 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
988 assert len(computed_blocks.blocks[0]) == 0
989 assert num_computed_tokens == 0
974 - block_hashes = manager.req_to_block_hashes[req2.request_id]
990 + block_hashes = req2.block_hashes
991 assert len(block_hashes) == 3
992 assert block_hashes[0].extra_keys == ("salt2", )
993
@@ -992,7 +1008,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1008 # Complete 3 blocks (48 tokens)
1009 # | Common-0 | Common-1 | Common-2 | ... |
1010 common_token_ids = [i for i in range(3) for _ in range(16)]
995 - req0 = make_request("0", common_token_ids)
1011 + req0 = make_request("0", common_token_ids, block_size, hash)
1012 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
1013 assert not computed_blocks.blocks[0]
1014 assert num_computed_tokens == 0
@@ -1003,7 +1019,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1019 req0.request_id]
1020
1021 # | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |
1006 - req1 = make_request("1", common_token_ids * 2)
1022 + req1 = make_request("1", common_token_ids * 2, block_size, hash)
1023 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
1024 assert computed_blocks.blocks[0] == block_part0
1025 assert num_computed_tokens == 3 * 16
@@ -1020,19 +1036,19 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1036
1037 # | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
1038 # | Req1-5(F)| Req2-0 | Req2-1 | ... |
1023 - req2 = make_request("2", [7] * block_size * 2)
1039 + req2 = make_request("2", [7] * block_size * 2, block_size, hash)
1040 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
1041 assert not computed_blocks.blocks[0]
1042 assert num_computed_tokens == 0
1043 manager.allocate_slots(req2, block_size * 2,
1028 - len(computed_blocks.blocks[0]) * 16,
1044 + len(computed_blocks.blocks[0]) * block_size,
1045 computed_blocks)
1046
1047 # Req3 is Req2 + 3 new blocks, so the first 6 blocks are computed,
1048 # but it cannot be allocated due to insufficient free blocks (2).
1049 # In this case, the ref_cnt of the computed blocks should not be changed.
1050 assert manager.block_pool.free_block_queue.num_free_blocks == 5
1035 - req3 = make_request("3", common_token_ids * 3)
1051 + req3 = make_request("3", common_token_ids * 3, block_size, hash)
1052 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
1053 assert computed_blocks.blocks[0] == block_part1
1054 assert num_computed_tokens == 6 * 16
@@ -1047,8 +1063,9 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1063
1064
1065 def test_reset_prefix_cache():
1066 + block_size = 16
1067 manager = KVCacheManager(
1051 - make_kv_cache_config(16, 11),
1068 + make_kv_cache_config(block_size, 11),
1069 max_model_len=8192,
1070 enable_caching=True,
1071 )
@@ -1056,15 +1073,15 @@ def test_reset_prefix_cache():
1073 full_block_token_ids = [i for i in range(3) for _ in range(16)]
1074 unique_token_ids = [3] * 7
1075 all_token_ids = full_block_token_ids + unique_token_ids
1059 - req0 = make_request("0", all_token_ids)
1076 + req0 = make_request("0", all_token_ids, block_size, hash)
1077 blocks = manager.allocate_slots(req0, 55)
1078 assert blocks.get_block_ids() == ([1, 2, 3, 4], )
1079
1080 unique_token_ids = [4] * 7
1081 all_token_ids = full_block_token_ids + unique_token_ids
1065 - req1 = make_request("1", all_token_ids)
1082 + req1 = make_request("1", all_token_ids, block_size, hash)
1083 computed_blocks, _ = manager.get_computed_blocks(req1)
1067 - assert len(manager.req_to_block_hashes[req1.request_id]) == 3
1084 + assert len(req1.block_hashes) == 3
1085 assert len(computed_blocks.blocks[0]) == 3
1086 blocks = manager.allocate_slots(req1, 7,
1087 len(computed_blocks.blocks[0]) * 16,
@@ -1086,8 +1103,9 @@ def test_reset_prefix_cache():
1103
1104 def test_prefix_cache_stats_disabled():
1105 """Test that prefix_cache_stats is None when log_stats is False."""
1106 + block_size = 16
1107 manager = KVCacheManager(
1090 - make_kv_cache_config(16, 11),
1108 + make_kv_cache_config(block_size, 11),
1109 max_model_len=8192,
1110 enable_caching=True,
1111 log_stats=False, # Disable logging stats
@@ -1095,7 +1113,7 @@ def test_prefix_cache_stats_disabled():
1113 assert manager.prefix_cache_stats is None
1114
1115 # Call all functions that check whether log_stats is disabled.
1098 - req = make_request("0", list(range(16)))
1116 + req = make_request("0", list(range(16)), block_size, hash)
1117 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
1118 assert not computed_blocks.blocks[0]
1119 assert num_computed_tokens == 0
@@ -1192,7 +1210,7 @@ def test_kv_cache_events(blocks_to_cache: int):
1210 )
1211
1212 num_tokens = block_size * blocks_to_cache
1195 - req0 = make_request("0", list(range(num_tokens)))
1213 + req0 = make_request("0", list(range(num_tokens)), block_size, hash)
1214 _ = manager.allocate_slots(req0, num_tokens)
1215 events = manager.take_events()
1216
@@ -1208,7 +1226,7 @@ def test_kv_cache_events(blocks_to_cache: int):
1226 # Should see block_to_cache number of removed block events and a new block
1227 # stored event
1228 manager.free(req0)
1211 - req1 = make_request("1", list(range(num_tokens)))
1229 + req1 = make_request("1", list(range(num_tokens)), block_size, hash)
1230 _ = manager.allocate_slots(req1, num_tokens)
1231 events = manager.take_events()
1232
@@ -1242,7 +1260,7 @@ def test_eagle_enabled_removes_last_block():
1260
1261 # Request with 3 full blocks (48 tokens)
1262 token_ids = [0] * (3 * block_size)
1245 - req = make_request("divisible_request", token_ids)
1263 + req = make_request("divisible_request", token_ids, block_size, hash)
1264
1265 # Prime the cache
1266 computed_blocks, _ = manager.get_computed_blocks(req)
@@ -1252,7 +1270,7 @@ def test_eagle_enabled_removes_last_block():
1270 manager.free(req)
1271
1272 # New request with same tokens + Eagle enabled
1255 - req_eagle = make_request("eagle_divisible", token_ids)
1273 + req_eagle = make_request("eagle_divisible", token_ids, block_size, hash)
1274 computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1275
1276 # Should retain 1 block:
@@ -1273,7 +1291,7 @@ def test_eagle_with_partial_blocks():
1291 )
1292 # 2 full blocks + 5 tokens (non-divisible length)
1293 token_ids = [0] * (2 * block_size + 5)
1276 - req = make_request("partial_block_test", token_ids)
1294 + req = make_request("partial_block_test", token_ids, block_size, hash)
1295
1296 # Prime the cache
1297 computed_blocks, _ = manager.get_computed_blocks(req)
@@ -1283,7 +1301,7 @@ def test_eagle_with_partial_blocks():
1301 manager.free(req)
1302
1303 # New request with Eagle enabled
1286 - req_eagle = make_request("partial_eagle", token_ids)
1304 + req_eagle = make_request("partial_eagle", token_ids, block_size, hash)
1305 computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1306 # Original match: 2 full blocks → Eagle removes 1 → 1 remaining
1307 assert len(computed_blocks.blocks[0]) == 1
@@ -1314,7 +1332,7 @@ def test_eagle_with_sliding_window():
1332
1333 # 2 full blocks + 5 tokens (non-divisible length)
1334 token_ids = [0] * (2 * block_size + 5)
1317 - req = make_request("partial_block_test", token_ids)
1335 + req = make_request("partial_block_test", token_ids, block_size, hash)
1336
1337 # Prime the cache
1338 computed_blocks, _ = manager.get_computed_blocks(req)
@@ -1322,12 +1340,12 @@ def test_eagle_with_sliding_window():
1340 len(computed_blocks.blocks[0]) * 16,
1341 computed_blocks)
1342 # record the block hash of the first block in the request for later use
1325 - block_hash_first_block = manager.req_to_block_hashes[req.request_id][0]
1343 + block_hash_first_block = req.block_hashes[0]
1344 assert block_hash_first_block is not None
1345 manager.free(req)
1346
1347 # New request with Eagle enabled
1330 - req_eagle = make_request("partial_eagle", token_ids)
1348 + req_eagle = make_request("partial_eagle", token_ids, block_size, hash)
1349 computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1350 # Original match: 2 full blocks → Eagle removes 1 → 1 remaining
1351 assert len(computed_blocks.blocks[0]) == 1
@@ -1340,7 +1358,8 @@ def test_eagle_with_sliding_window():
1358 BlockHashWithGroupId(block_hash_first_block, 0))
1359
1360 # New request
1343 - req_after_evict = make_request("partial_eagle_after_evict", token_ids)
1361 + req_after_evict = make_request("partial_eagle_after_evict", token_ids,
1362 + block_size, hash)
1363 computed_blocks, num_tokens = manager.get_computed_blocks(req_after_evict)
1364 # Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is
1365 # not considered. But after dropping the last matched block due to eagle,
tests/v1/core/test_scheduler.py
+16 -13
@@ -589,7 +589,7 @@ def test_preempt_during_execution():
589 block_size=16,
590 num_blocks=11,
591 enable_prefix_caching=False)
592 - requests = create_requests(num_requests=2, num_tokens=80)
592 + requests = create_requests(num_requests=2, num_tokens=80, block_size=16)
593
594 # Schedule the first request.
595 scheduler.add_request(requests[0])
@@ -762,7 +762,7 @@ def _assert_right_scheduler_output(
762
763 def _assert_right_kv_cache_manager(
764 scheduler: Scheduler,
765 - req_ids: list[str],
765 + requests: list[Request],
766 num_tokens: int,
767 block_size: int,
768 num_requests: int,
@@ -772,12 +772,12 @@ def _assert_right_kv_cache_manager(
772
773 # Make sure the request stats are right.
774 EXPECTED_TOTAL_BLOCKS = num_tokens // block_size
775 - for req_id in req_ids:
775 + for req in requests:
776 blocks = (scheduler.kv_cache_manager.coordinator.
777 - single_type_managers[0].req_to_blocks[req_id])
778 - hashes = scheduler.kv_cache_manager.req_to_block_hashes[req_id]
777 + single_type_managers[0].req_to_blocks[req.request_id])
778 + hashes = req.block_hashes
779 assert (scheduler.kv_cache_manager.coordinator.single_type_managers[0].
780 - num_cached_block[req_id] == EXPECTED_TOTAL_BLOCKS)
780 + num_cached_block[req.request_id] == EXPECTED_TOTAL_BLOCKS)
781 assert len(blocks) == EXPECTED_TOTAL_BLOCKS
782 assert len(hashes) == EXPECTED_TOTAL_BLOCKS
783
@@ -840,7 +840,8 @@ def test_kv_connector_basic():
840 MAX_TOKENS = 3
841 requests = create_requests(num_requests=NUM_REQUESTS,
842 num_tokens=NUM_TOKENS,
843 - max_tokens=MAX_TOKENS)
843 + max_tokens=MAX_TOKENS,
844 + block_size=BLOCK_SIZE)
845 req_ids = []
846 req_to_index = {}
847 for i, request in enumerate(requests):
@@ -868,7 +869,7 @@ def test_kv_connector_basic():
869 )
870
871 # Ensure KVCacheManager is correct.
871 - _assert_right_kv_cache_manager(scheduler, req_ids, NUM_TOKENS, BLOCK_SIZE,
872 + _assert_right_kv_cache_manager(scheduler, requests, NUM_TOKENS, BLOCK_SIZE,
873 NUM_REQUESTS, NUM_TOTAL_BLOCKS)
874
875 # Continue Generation until done.
@@ -886,7 +887,8 @@ def test_kv_connector_basic():
887 NUM_TOKENS = NUM_TOKENS_PREFIX * 2
888 requests = create_requests(num_requests=NUM_REQUESTS,
889 num_tokens=NUM_TOKENS,
889 - max_tokens=MAX_TOKENS)
890 + max_tokens=MAX_TOKENS,
891 + block_size=BLOCK_SIZE)
892 req_ids = []
893 req_to_index = {}
894 for i, request in enumerate(requests):
@@ -915,7 +917,7 @@ def test_kv_connector_basic():
917 NUM_MATCHED_NEW_TOKENS))
918
919 # Ensure KVCacheManager is correct.
918 - _assert_right_kv_cache_manager(scheduler, req_ids, NUM_TOKENS, BLOCK_SIZE,
920 + _assert_right_kv_cache_manager(scheduler, requests, NUM_TOKENS, BLOCK_SIZE,
921 NUM_REQUESTS, NUM_TOTAL_BLOCKS)
922
923 # Continue Generation until done.
@@ -953,7 +955,8 @@ def test_kv_connector_unable_to_allocate():
955 MAX_TOKENS = 2
956 requests = create_requests(num_requests=NUM_REQUESTS,
957 num_tokens=NUM_TOKENS,
956 - max_tokens=MAX_TOKENS)
958 + max_tokens=MAX_TOKENS,
959 + block_size=BLOCK_SIZE)
960 req_ids = []
961 req_to_index = {}
962 for i, request in enumerate(requests):
@@ -1034,7 +1037,8 @@ def test_kv_connector_handles_preemption():
1037 MAX_TOKENS = BLOCK_SIZE * 2
1038 requests = create_requests(num_requests=NUM_REQUESTS,
1039 num_tokens=NUM_TOKENS,
1037 - max_tokens=MAX_TOKENS)
1040 + max_tokens=MAX_TOKENS,
1041 + block_size=BLOCK_SIZE)
1042 req_ids = []
1043 req_to_index = {}
1044 for i, request in enumerate(requests):
@@ -1162,7 +1166,6 @@ def assert_scheduler_empty(scheduler: Scheduler):
1166 # KVCache Manager.
1167 assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
1168 req_to_blocks) == 0
1165 - assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
1169 assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
1170 num_cached_block) == 0
1171 num_free_blocks = (
tests/v1/core/test_single_type_kv_cache_manager.py
-2
@@ -17,7 +17,6 @@ from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
17 def get_sliding_window_manager(sliding_window_spec, block_pool):
18 return SlidingWindowManager(sliding_window_spec,
19 block_pool,
20 - caching_hash_fn=lambda x: x,
20 kv_cache_group_id=0)
21
22
@@ -25,7 +24,6 @@ def get_chunked_local_attention_manager(chunked_local_attention_spec,
24 block_pool):
25 return ChunkedLocalAttentionManager(chunked_local_attention_spec,
26 block_pool,
28 - caching_hash_fn=lambda x: x,
27 kv_cache_group_id=0)
28
29
tests/v1/core/utils.py
+16 -1
@@ -10,6 +10,8 @@ from vllm.multimodal.inputs import (MultiModalBatchedField,
10 MultiModalFieldElem, MultiModalKwargsItem,
11 PlaceholderRange)
12 from vllm.sampling_params import SamplingParams
13 +from vllm.v1.core.kv_cache_utils import (get_request_block_hasher,
14 + init_none_hash)
15 from vllm.v1.core.sched.async_scheduler import AsyncScheduler
16 from vllm.v1.core.sched.scheduler import Scheduler
17 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
@@ -114,6 +116,9 @@ def create_scheduler(
116 )
117
118
119 +_none_hash_initialized = False
120 +
121 +
122 def create_requests(
123 num_requests: int,
124 num_tokens: int = 10,
@@ -122,7 +127,14 @@ def create_requests(
127 stop_token_ids: Optional[list[int]] = None,
128 prompt_logprobs: Optional[int] = None,
129 same_prompt: bool = False,
130 + block_size: int = 16,
131 ) -> list[Request]:
132 + global _none_hash_initialized
133 + if not _none_hash_initialized:
134 + init_none_hash(hash)
135 + _none_hash_initialized = True
136 +
137 + block_hasher = get_request_block_hasher(block_size, hash)
138 sampling_params = SamplingParams(ignore_eos=False,
139 max_tokens=max_tokens,
140 stop_token_ids=stop_token_ids,
@@ -139,9 +151,11 @@ def create_requests(
151 )
152 mm_item = MultiModalKwargsItem.from_elems([mm_elem])
153 mm_kwargs = [mm_item] * len(mm_position)
154 + mm_hashes = ["hash"] * len(mm_position)
155 else:
156 mm_position = None
157 mm_kwargs = None
158 + mm_hashes = None
159 prompt_token_ids = ([0] * num_tokens if same_prompt else [i] *
160 num_tokens)
161 request = Request(
@@ -151,8 +165,9 @@ def create_requests(
165 pooling_params=None,
166 multi_modal_kwargs=mm_kwargs,
167 multi_modal_placeholders=mm_position,
154 - multi_modal_hashes=None,
168 + multi_modal_hashes=mm_hashes,
169 eos_token_id=EOS_TOKEN_ID,
170 + block_hasher=block_hasher,
171 )
172 requests.append(request)
173 return requests
tests/v1/kv_connector/unit/test_nixl_connector.py
+2
@@ -147,6 +147,7 @@ def test_basic_interface():
147 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
148
149 request = create_request(request_id=1,
150 + block_size=BLOCK_SIZE,
151 num_tokens=NUM_TOKENS,
152 do_remote_prefill=True)
153 request_id = request.request_id
@@ -186,6 +187,7 @@ def test_prompt_less_than_block_size():
187
188 # Request will have 1 partial remote block.
189 request = create_request(request_id=1,
190 + block_size=BLOCK_SIZE,
191 num_tokens=NUM_TOKENS,
192 do_remote_prefill=True,
193 num_remote_blocks=1)
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+8 -2
@@ -21,6 +21,7 @@ def test_basic_lifecycle():
21 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
22
23 request = create_request(request_id=1,
24 + block_size=BLOCK_SIZE,
25 max_tokens=1,
26 num_tokens=NUM_TOKENS,
27 do_remote_decode=True)
@@ -103,8 +104,10 @@ def test_short_prompt_lifecycle():
104 scheduler = create_scheduler(vllm_config)
105
106 # Not enough tokens for full block.
106 - NUM_TOKENS = vllm_config.cache_config.block_size // 2
107 + BLOCK_SIZE = vllm_config.cache_config.block_size
108 + NUM_TOKENS = BLOCK_SIZE // 2
109 request = create_request(request_id=1,
110 + block_size=BLOCK_SIZE,
111 max_tokens=1,
112 num_tokens=NUM_TOKENS,
113 do_remote_decode=True)
@@ -148,7 +151,9 @@ def test_prefix_cache_lifecycle():
151 NUM_EXTERNAL_FULL_BLOCKS = 3
152 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
153
151 - request_normal = create_request(request_id=1, num_tokens=NUM_TOKENS)
154 + request_normal = create_request(request_id=1,
155 + block_size=BLOCK_SIZE,
156 + num_tokens=NUM_TOKENS)
157
158 scheduler.add_request(request_normal)
159 scheduler_output = scheduler.schedule()
@@ -166,6 +171,7 @@ def test_prefix_cache_lifecycle():
171 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
172
173 request_remote = create_request(request_id=1,
174 + block_size=BLOCK_SIZE,
175 num_tokens=NUM_TOKENS,
176 do_remote_decode=True)
177
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+15 -2
@@ -23,6 +23,7 @@ def test_basic_lifecycle():
23 scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
24
25 request = create_request(request_id=1,
26 + block_size=BLOCK_SIZE,
27 num_tokens=NUM_TOKENS,
28 do_remote_prefill=True)
29
@@ -133,14 +134,17 @@ def test_interleaved_lifecycle():
134 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
135
136 request_remote = create_request(request_id=1,
137 + block_size=BLOCK_SIZE,
138 num_tokens=NUM_TOKENS,
139 do_remote_prefill=True)
140 request_local_a = create_request(
141 request_id=2,
142 + block_size=BLOCK_SIZE,
143 num_tokens=NUM_TOKENS,
144 )
145 request_local_b = create_request(
146 request_id=3,
147 + block_size=BLOCK_SIZE,
148 num_tokens=NUM_TOKENS,
149 )
150
@@ -236,6 +240,7 @@ def test_no_spurious_prefix_caching():
240 # Both of these requests have prompts like [1,1,1,1,1, ...]
241 request_remote = create_request(
242 request_id=1,
243 + block_size=BLOCK_SIZE,
244 num_tokens=NUM_TOKENS,
245 do_remote_prefill=True,
246 use_all_1s_for_prompt_tokens=True,
@@ -243,6 +248,7 @@ def test_no_spurious_prefix_caching():
248
249 request_local = create_request(
250 request_id=2,
251 + block_size=BLOCK_SIZE,
252 num_tokens=NUM_TOKENS,
253 do_remote_prefill=False,
254 use_all_1s_for_prompt_tokens=True,
@@ -292,6 +298,7 @@ def test_full_block_prompt():
298 NUM_TOKENS = int(BLOCK_SIZE * NUM_EXTERNAL_FULL_BLOCKS)
299
300 request = create_request(request_id=1,
301 + block_size=BLOCK_SIZE,
302 num_tokens=NUM_TOKENS,
303 do_remote_prefill=True)
304
@@ -364,8 +371,11 @@ def test_cannot_schedule_after_recv():
371 NUM_TOKENS_LOCAL = int(BLOCK_SIZE * NUM_PROMPT_BLOCKS)
372 NUM_TOKENS_REMOTE = int(BLOCK_SIZE * NUM_PROMPT_BLOCKS)
373
367 - request_normal = create_request(request_id=1, num_tokens=NUM_TOKENS_LOCAL)
374 + request_normal = create_request(request_id=1,
375 + block_size=BLOCK_SIZE,
376 + num_tokens=NUM_TOKENS_LOCAL)
377 request_remote = create_request(request_id=2,
378 + block_size=BLOCK_SIZE,
379 num_tokens=NUM_TOKENS_REMOTE,
380 do_remote_prefill=True)
381
@@ -456,8 +466,11 @@ def test_cannot_recv():
466 NUM_TOKENS_LOCAL = int(BLOCK_SIZE * NUM_PROMPT_BLOCKS)
467 NUM_TOKENS_REMOTE = int(BLOCK_SIZE * (NUM_PROMPT_BLOCKS + 0.5))
468
459 - request_normal = create_request(request_id=1, num_tokens=NUM_TOKENS_LOCAL)
469 + request_normal = create_request(request_id=1,
470 + block_size=BLOCK_SIZE,
471 + num_tokens=NUM_TOKENS_LOCAL)
472 request_remote = create_request(request_id=2,
473 + block_size=BLOCK_SIZE,
474 num_tokens=NUM_TOKENS_REMOTE,
475 do_remote_prefill=True)
476
tests/v1/kv_connector/unit/utils.py
+20 -11
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import tempfile
4 from collections import defaultdict
5 -from typing import Any, Optional
5 +from typing import Any, Callable, Optional
6
7 import torch
8
@@ -14,6 +14,8 @@ from vllm.distributed.kv_transfer.kv_connector.factory import (
14 from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa
15 SharedStorageConnector)
16 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17 +from vllm.v1.core.kv_cache_utils import (get_request_block_hasher,
18 + init_none_hash)
19 from vllm.v1.core.sched.scheduler import Scheduler
20 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
21 KVCacheGroupSpec)
@@ -40,7 +42,6 @@ def assert_scheduler_empty(scheduler: Scheduler):
42 # KVCache Manager.
43 assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
44 req_to_blocks) == 0
43 - assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
45 assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
46 num_cached_block) == 0
47 num_free_blocks = (
@@ -115,16 +116,23 @@ def create_scheduler(
116 )
117
118
118 -def create_request(
119 - request_id: int,
120 - num_tokens: int = 10,
121 - max_tokens: int = 16,
122 - do_remote_decode: bool = False,
123 - do_remote_prefill: bool = False,
124 - use_all_1s_for_prompt_tokens: bool = False,
125 - num_remote_blocks: int = 3,
126 -) -> Request:
119 +_none_hash_initialized = False
120 +
121 +
122 +def create_request(request_id: int,
123 + num_tokens: int = 10,
124 + max_tokens: int = 16,
125 + do_remote_decode: bool = False,
126 + do_remote_prefill: bool = False,
127 + use_all_1s_for_prompt_tokens: bool = False,
128 + num_remote_blocks: int = 3,
129 + block_size: int = 16,
130 + hash_fn: Callable = hash) -> Request:
131 """Make dummy request for testing."""
132 + global _none_hash_initialized
133 + if not _none_hash_initialized:
134 + init_none_hash(hash)
135 + _none_hash_initialized = True
136
137 kv_transfer_params: Optional[dict[str, Any]] = None
138
@@ -158,6 +166,7 @@ def create_request(
166 multi_modal_placeholders=None,
167 multi_modal_hashes=None,
168 eos_token_id=EOS_TOKEN_ID,
169 + block_hasher=get_request_block_hasher(block_size, hash_fn),
170 )
171 req.kv_transfer_params = kv_transfer_params
172 return req
vllm/utils/__init__.py
+18
@@ -3243,6 +3243,24 @@ def sha256_cbor_64bit(input) -> int:
3243 return full_hash & ((1 << 64) - 1)
3244
3245
3246 +def get_hash_fn_by_name(hash_fn_name: str) -> Callable:
3247 + """Get a hash function by name, or raise an error if
3248 + the function is not found.
3249 + Args:
3250 + hash_fn_name: Name of the hash function.
3251 + Returns:
3252 + A hash function.
3253 + """
3254 + if hash_fn_name == "sha256":
3255 + return sha256
3256 + if hash_fn_name == "sha256_cbor_64bit":
3257 + return sha256_cbor_64bit
3258 + if hash_fn_name == "builtin":
3259 + return hash
3260 +
3261 + raise ValueError(f"Unsupported hash function: {hash_fn_name}")
3262 +
3263 +
3264 def is_torch_equal_or_newer(target: str) -> bool:
3265 """Check if the installed torch version is >= the target version.
3266
vllm/v1/core/block_pool.py
+17 -58
@@ -2,15 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import defaultdict
4 from collections.abc import Iterable
5 -from typing import Callable, Optional
5 +from typing import Optional
6
7 from vllm.distributed.kv_events import (AllBlocksCleared, BlockRemoved,
8 BlockStored, KVCacheEvent)
9 from vllm.logger import init_logger
10 from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
11 - FreeKVCacheBlockQueue, KVCacheBlock,
12 - generate_block_hash_extra_keys,
13 - hash_block_tokens)
11 + FreeKVCacheBlockQueue, KVCacheBlock)
12 from vllm.v1.request import Request
13
14 logger = init_logger(__name__)
@@ -97,84 +95,39 @@ class BlockPool:
95 self,
96 request: Request,
97 blocks: list[KVCacheBlock],
100 - block_hashes: list[BlockHash],
98 num_cached_blocks: int,
99 num_full_blocks: int,
100 block_size: int,
101 kv_cache_group_id: int,
105 - hash_fn: Callable,
102 ) -> None:
103 """Cache a list of full blocks for prefix caching.
104 This function takes a list of blocks that will have their block hash
109 - metadata to be updated and cached. Given a request, it computes the
110 - block hashes for the blocks starting from `num_cached_blocks` to
111 - `num_full_blocks`, updating the metadata for each block
112 - and caching them in the `cached_block_hash_to_block`.
105 + metadata to be updated and cached. Given a request, it updates the
106 + metadata for each block and caching it in the
107 + `cached_block_hash_to_block`.
108 + The block hashes values are computed by the Request object immediately
109 + when it is created and when new tokens are appended.
110
111 Args:
112 request: The request to cache the blocks.
113 blocks: All blocks in the request.
117 - block_hashes: Block hashes of the blocks in the request. Note that
118 - this list may be shorter than the blocks list. In this case the
119 - missed block hash will be computed in this function.
114 num_cached_blocks: The number of blocks that are already cached.
115 num_full_blocks: The number of blocks that are full and should
116 be cached after this function.
117 block_size: Number of tokens in each block.
118 kv_cache_group_id: The id of the KV cache group.
125 - hash_fn: The hash function to use for block hashes.
119 """
120 if num_cached_blocks == num_full_blocks:
121 return
122 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
130 - assert len(block_hashes) >= num_cached_blocks
131 - new_block_hashes = block_hashes[num_cached_blocks:]
123 + assert len(request.block_hashes) >= num_full_blocks
124 + new_block_hashes = request.block_hashes[num_cached_blocks:]
125
133 - # Update the new blocks with the block hashes through the chain.
134 - if num_cached_blocks == 0:
135 - prev_block_hash_value = None
136 - else:
137 - prev_block = blocks[num_cached_blocks - 1]
138 - assert prev_block.block_hash is not None
139 - prev_block_hash_value = prev_block.block_hash.get_hash_value()
140 -
141 - parent_block_hash = prev_block_hash_value
126 new_hashes: Optional[list[int]] = ([] if self.enable_kv_cache_events
127 else None)
128 for i, blk in enumerate(new_full_blocks):
129 assert blk.block_hash is None
146 -
147 - if i < len(new_block_hashes):
148 - # The block hash may already be computed in
149 - # "get_computed_blocks" if the tokens are not generated by
150 - # this request (either the prompt tokens or the previously
151 - # generated tokens with preemption), or by other
152 - # single_type_managers with the same block_size.
153 - # In this case we simply reuse the block hash.
154 - block_hash = new_block_hashes[i]
155 - else:
156 - # Otherwise compute the block hash and cache it in the request
157 - # in case it will be preempted in the future.
158 - blk_idx = num_cached_blocks + i
159 - start_token_idx = blk_idx * block_size
160 - end_token_idx = (blk_idx + 1) * block_size
161 - block_tokens = request.all_token_ids[
162 - start_token_idx:end_token_idx]
163 - assert len(block_tokens) == block_size, (
164 - f"Expected {block_size} tokens, got "
165 - f"{len(block_tokens)} at {blk_idx}th block for request "
166 - f"{request.request_id}({request})")
167 -
168 - # Generate extra keys for multi-modal inputs. Note that since
169 - # we reach to this branch only when the block is completed with
170 - # generated tokens, we only need to consider the last mm input.
171 - extra_keys, _ = generate_block_hash_extra_keys(
172 - request, start_token_idx, end_token_idx, -1)
173 -
174 - # Compute the hash of the current block.
175 - block_hash = hash_block_tokens(hash_fn, prev_block_hash_value,
176 - block_tokens, extra_keys)
177 - block_hashes.append(block_hash)
130 + block_hash = new_block_hashes[i]
131
132 # Update and added the full block to the cache.
133 block_hash_with_group_id = BlockHashWithGroupId(
@@ -184,9 +137,15 @@ class BlockPool:
137 blk.block_id] = blk
138 if new_hashes is not None:
139 new_hashes.append(block_hash.hash_value)
187 - prev_block_hash_value = block_hash.hash_value
140
141 if self.enable_kv_cache_events:
142 + if num_cached_blocks == 0:
143 + parent_block_hash = None
144 + else:
145 + parent_block = blocks[num_cached_blocks - 1]
146 + assert parent_block.block_hash is not None
147 + parent_block_hash = parent_block.block_hash.get_hash_value()
148 +
149 self.kv_event_queue.append(
150 BlockStored(
151 block_hashes=new_hashes,
vllm/v1/core/kv_cache_coordinator.py
+12 -21
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 -from typing import Callable, Optional
4 +from typing import Optional
5
6 from vllm.v1.core.block_pool import BlockPool
7 from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
@@ -23,7 +23,6 @@ class KVCacheCoordinator(ABC):
23 max_model_len: int,
24 use_eagle: bool,
25 enable_caching: bool,
26 - caching_hash_fn: Callable,
26 enable_kv_cache_events: bool,
27 ):
28 self.kv_cache_config = kv_cache_config
@@ -40,7 +39,6 @@ class KVCacheCoordinator(ABC):
39 kv_cache_spec=kv_cache_group.kv_cache_spec,
40 block_pool=self.block_pool,
41 kv_cache_group_id=i,
43 - caching_hash_fn=caching_hash_fn,
42 ) for i, kv_cache_group in enumerate(
43 self.kv_cache_config.kv_cache_groups))
44
@@ -99,19 +97,17 @@ class KVCacheCoordinator(ABC):
97 manager.allocate_new_blocks(request_id, num_tokens)
98 for manager in self.single_type_managers)
99
102 - def cache_blocks(self, request: Request, block_hashes: list[BlockHash],
103 - num_computed_tokens: int) -> None:
100 + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
101 """
102 Cache the blocks for the request.
103
104 Args:
105 request: The request.
109 - block_hashes: The block hashes of the request.
106 num_tokens: The total number of tokens that need to be cached
107 (including tokens that are already cached).
108 """
109 for manager in self.single_type_managers:
114 - manager.cache_blocks(request, block_hashes, num_computed_tokens)
110 + manager.cache_blocks(request, num_computed_tokens)
111
112 def free(self, request_id: str) -> None:
113 """
@@ -184,10 +180,9 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
180 """
181
182 def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
187 - use_eagle: bool, caching_hash_fn: Callable,
188 - enable_kv_cache_events: bool):
183 + use_eagle: bool, enable_kv_cache_events: bool):
184 super().__init__(kv_cache_config, max_model_len, use_eagle, False,
190 - caching_hash_fn, enable_kv_cache_events)
185 + enable_kv_cache_events)
186 self.num_single_type_manager = len(self.single_type_managers)
187
188 def get_num_common_prefix_blocks(self, request_id: str,
@@ -213,10 +208,9 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
208
209 def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
210 use_eagle: bool, enable_caching: bool,
216 - caching_hash_fn: Callable, enable_kv_cache_events: bool):
211 + enable_kv_cache_events: bool):
212 super().__init__(kv_cache_config, max_model_len, use_eagle,
218 - enable_caching, caching_hash_fn,
219 - enable_kv_cache_events)
213 + enable_caching, enable_kv_cache_events)
214 self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[
215 0].kv_cache_spec
216 self.block_size = self.kv_cache_spec.block_size
@@ -250,10 +244,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
244
245 def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
246 use_eagle: bool, enable_caching: bool,
253 - caching_hash_fn: Callable, enable_kv_cache_events: bool):
247 + enable_kv_cache_events: bool):
248 super().__init__(kv_cache_config, max_model_len, use_eagle,
255 - enable_caching, caching_hash_fn,
256 - enable_kv_cache_events)
249 + enable_caching, enable_kv_cache_events)
250 self.verify_and_split_kv_cache_groups()
251
252 def verify_and_split_kv_cache_groups(self) -> None:
@@ -386,17 +379,15 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
379
380 def get_kv_cache_coordinator(
381 kv_cache_config: KVCacheConfig, max_model_len: int, use_eagle: bool,
389 - enable_caching: bool, caching_hash_fn: Callable,
382 + enable_caching: bool,
383 enable_kv_cache_events: bool) -> KVCacheCoordinator:
384 if not enable_caching:
385 return KVCacheCoordinatorNoPrefixCache(kv_cache_config, max_model_len,
393 - use_eagle, caching_hash_fn,
386 + use_eagle,
387 enable_kv_cache_events)
388 if len(kv_cache_config.kv_cache_groups) == 1:
389 return UnitaryKVCacheCoordinator(kv_cache_config, max_model_len,
390 use_eagle, enable_caching,
398 - caching_hash_fn,
391 enable_kv_cache_events)
392 return HybridKVCacheCoordinator(kv_cache_config, max_model_len, use_eagle,
401 - enable_caching, caching_hash_fn,
402 - enable_kv_cache_events)
393 + enable_caching, enable_kv_cache_events)
vllm/v1/core/kv_cache_manager.py
+4 -47
@@ -1,16 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from collections import defaultdict
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 sha256, sha256_cbor_64bit
9 from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
12 -from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
13 - hash_request_tokens, init_none_hash)
10 +from vllm.v1.core.kv_cache_utils import KVCacheBlock
11 from vllm.v1.kv_cache_interface import KVCacheConfig
12 from vllm.v1.metrics.stats import PrefixCacheStats
13 from vllm.v1.request import Request, RequestStatus
@@ -71,23 +68,13 @@ class KVCacheManager:
68 kv_cache_config: KVCacheConfig,
69 max_model_len: int,
70 enable_caching: bool = True,
74 - caching_hash_algo: str = "builtin",
71 use_eagle: bool = False,
72 log_stats: bool = False,
73 enable_kv_cache_events: bool = False,
74 ) -> None:
75 self.max_model_len = max_model_len
76
81 - if len(kv_cache_config.kv_cache_groups) == 0:
82 - # Attention free models don't have kv cache,
83 - # thus don't need prefix caching.
84 - enable_caching = False
77 self.enable_caching = enable_caching
86 -
87 - self.caching_hash_fn = (
88 - sha256_cbor_64bit if caching_hash_algo == "sha256_cbor_64bit" else
89 - sha256 if caching_hash_algo == "sha256" else hash)
90 - init_none_hash(self.caching_hash_fn)
78 self.use_eagle = use_eagle
79 self.log_stats = log_stats
80 # FIXME: make prefix cache stats conditional on log_stats
@@ -107,19 +94,12 @@ class KVCacheManager:
94 max_model_len=self.max_model_len,
95 use_eagle=self.use_eagle,
96 enable_caching=self.enable_caching,
110 - caching_hash_fn=self.caching_hash_fn,
97 enable_kv_cache_events=enable_kv_cache_events,
98 )
99 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
100 self.block_pool = self.coordinator.block_pool
101 self.kv_cache_config = kv_cache_config
102
117 - # Mapping from request ID to kv block hashes.
118 - # This is to avoid recomputing the block hashes for each call of
119 - # `get_computed_blocks` or `allocate_slots`.
120 - self.req_to_block_hashes: defaultdict[
121 - str, list[BlockHash]] = defaultdict(list)
122 -
103 @property
104 def usage(self) -> float:
105 """Get the KV cache usage.
@@ -161,15 +141,6 @@ class KVCacheManager:
141 and request.sampling_params.prompt_logprobs is not None)):
142 return self.create_empty_block_list(), 0
143
164 - # The block hashes for the request may already be computed
165 - # if the scheduler has tried to schedule the request before.
166 - block_hashes = self.req_to_block_hashes[request.request_id]
167 - if not block_hashes:
168 - assert self.block_size is not None
169 - block_hashes = hash_request_tokens(self.caching_hash_fn,
170 - self.block_size, request)
171 - self.req_to_block_hashes[request.request_id] = block_hashes
172 -
144 # NOTE: When all tokens hit the cache, we must recompute the last token
145 # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
146 # This can trigger recomputation of an entire block, rather than just
@@ -178,7 +149,7 @@ class KVCacheManager:
149 # could slightly improve performance in the future.
150 max_cache_hit_length = request.num_tokens - 1
151 computed_blocks, num_new_computed_tokens = (
181 - self.coordinator.find_longest_cache_hit(block_hashes,
152 + self.coordinator.find_longest_cache_hit(request.block_hashes,
153 max_cache_hit_length))
154
155 if self.log_stats:
@@ -296,11 +267,7 @@ class KVCacheManager:
267 # at `request.num_tokens`, ensuring only "finalized" tokens are cached.
268 num_tokens_to_cache = min(num_computed_tokens + num_new_tokens,
269 request.num_tokens)
299 - self.coordinator.cache_blocks(
300 - request,
301 - self.req_to_block_hashes[request.request_id],
302 - num_tokens_to_cache,
303 - )
270 + self.coordinator.cache_blocks(request, num_tokens_to_cache)
271
272 return KVCacheBlocks(new_blocks)
273
@@ -373,14 +340,6 @@ class KVCacheManager:
340 return self.coordinator.get_num_common_prefix_blocks(
341 request.request_id, num_running_requests)
342
376 - def free_block_hashes(self, request: Request) -> None:
377 - """Discard the block hashes for the request.
378 -
379 - NOTE: Unlike `free`, this method should be called only when the request
380 - is finished, not when it is preempted.
381 - """
382 - self.req_to_block_hashes.pop(request.request_id, None)
383 -
343 def take_events(self) -> list[KVCacheEvent]:
344 """Take the KV cache events from the block pool.
345
@@ -397,9 +356,7 @@ class KVCacheManager:
356 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
357 """Cache the blocks for the request, if enabled."""
358 if self.enable_caching:
400 - block_hashes = self.req_to_block_hashes[request.request_id]
401 - self.coordinator.cache_blocks(request, block_hashes,
402 - num_computed_tokens)
359 + self.coordinator.cache_blocks(request, num_computed_tokens)
360
361 def create_empty_block_list(self) -> KVCacheBlocks:
362 """Creates a new KVCacheBlocks instance with no blocks."""
vllm/v1/core/kv_cache_utils.py
+50 -30
@@ -547,41 +547,61 @@ def hash_block_tokens(
547 curr_block_token_ids_tuple, extra_keys)
548
549
550 -def hash_request_tokens(hash_function: Any, block_size: int,
551 - request: Request) -> list[BlockHash]:
552 - """Computes hash values of a chain of blocks given a sequence of
553 - token IDs. The hash value is used for prefix caching.
550 +def get_request_block_hasher(
551 + block_size: int,
552 + caching_hash_fn: Callable[[Any],
553 + int]) -> Callable[[Request], list[BlockHash]]:
554 + """
555 + Returns a function which computes the list of un-computed block hashes
556 + of a request.
557 +
558 + Each request holds a list of its block hashes (request.block_hashes).
559 + When a request is created, it calls the below function to compute
560 + the hashes of all full blocks of the request's initial tokens.
561 + The hashes are then stored in request.block_hashes.
562 + Later, whenever new tokens are appended to the request, it calls
563 + the below function again to compute any new full blocks of tokens.
564 + The returned new hashes are appended to request.block_hashes.
565 + """
566
555 - Args:
556 - block_size: The size of each block.
557 - request: The request object.
567 + def request_block_hasher(request: Request) -> list[BlockHash]:
568 + start_token_idx = len(request.block_hashes) * block_size
569 + num_tokens = request.num_tokens
570 +
571 + curr_mm_idx = 0
572 + if start_token_idx > 0:
573 + # Set curr_mm_idx = -1 to indicate the last mm input.
574 + # Note that since we reach to this branch only when the block is
575 + # completed with generated tokens, we only need to consider the
576 + # last mm input.
577 + curr_mm_idx = -1
578 +
579 + prev_block_hash_value = request.block_hashes[-1].hash_value \
580 + if request.block_hashes else None
581 + new_block_hashes: list[BlockHash] = []
582 + while True:
583 + end_token_idx = start_token_idx + block_size
584 + if end_token_idx > num_tokens:
585 + # We only hash full blocks
586 + break
587
559 - Returns:
560 - The list of computed hash values.
561 - """
562 - token_ids = request.all_token_ids
588 + # MM and LoRA requests need extra keys for block-hash computation.
589 + extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
590 + request, start_token_idx, end_token_idx, curr_mm_idx)
591
564 - req_need_extra_keys = need_extra_keys(request)
565 - req_extra_keys = None
566 - curr_mm_idx = 0
592 + # Compute the hash of the current block
593 + block_tokens = request.all_token_ids[start_token_idx:end_token_idx]
594 + block_hash = hash_block_tokens(caching_hash_fn,
595 + prev_block_hash_value, block_tokens,
596 + extra_keys)
597
568 - ret = []
569 - parent_block_hash_value = None
570 - # Only full blocks will be hashed
571 - for start in range(0, len(token_ids) - block_size + 1, block_size):
572 - end = start + block_size
573 - block_token_ids = token_ids[start:end]
598 + new_block_hashes.append(block_hash)
599 + start_token_idx += block_size
600 + prev_block_hash_value = block_hash.hash_value
601
575 - if req_need_extra_keys:
576 - # MM and LoRA requests need extra keys for block-hash computation.
577 - req_extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
578 - request, start, end, curr_mm_idx)
579 -
580 - block_hash = hash_block_tokens(hash_function, parent_block_hash_value,
581 - block_token_ids, req_extra_keys)
582 - ret.append(block_hash)
583 - parent_block_hash_value = block_hash.hash_value
584 - return ret
602 + return new_block_hashes
603 +
604 + return request_block_hasher
605
606
607 def max_memory_usage_bytes(vllm_config: VllmConfig,
vllm/v1/core/sched/scheduler.py
-2
@@ -155,7 +155,6 @@ class Scheduler(SchedulerInterface):
155 kv_cache_config=kv_cache_config,
156 max_model_len=self.max_model_len,
157 enable_caching=self.cache_config.enable_prefix_caching,
158 - caching_hash_algo=self.cache_config.prefix_caching_hash_algo,
158 use_eagle=self.use_eagle,
159 log_stats=self.log_stats,
160 enable_kv_cache_events=self.enable_kv_cache_events,
@@ -1036,7 +1035,6 @@ class Scheduler(SchedulerInterface):
1035 def _free_blocks(self, request: Request):
1036 assert request.is_finished()
1037 self.kv_cache_manager.free(request)
1039 - self.kv_cache_manager.free_block_hashes(request)
1038 del self.requests[request.request_id]
1039
1040 def get_num_unfinished_requests(self) -> int:
vllm/v1/core/single_type_kv_cache_manager.py
+1 -9
@@ -3,7 +3,6 @@
3 import itertools
4 from abc import ABC, abstractmethod
5 from collections import defaultdict
6 -from typing import Callable
6
7 from vllm.utils import cdiv
8 from vllm.v1.core.block_pool import BlockPool
@@ -25,7 +24,6 @@ class SingleTypeKVCacheManager(ABC):
24 kv_cache_spec: KVCacheSpec,
25 block_pool: BlockPool,
26 kv_cache_group_id: int,
28 - caching_hash_fn: Callable,
27 ) -> None:
28 """
29 Initializes the SingleTypeKVCacheManager.
@@ -33,7 +31,6 @@ class SingleTypeKVCacheManager(ABC):
31 kv_cache_spec: The kv_cache_spec for this manager.
32 block_pool: The block pool.
33 kv_cache_group_id: The id of the kv cache group of this manager.
36 - caching_hash_fn: The caching hash function.
34 """
35
36 self.block_size = kv_cache_spec.block_size
@@ -52,7 +49,6 @@ class SingleTypeKVCacheManager(ABC):
49 # data for reempted ones.
50 self.num_cached_block: dict[str, int] = {}
51
55 - self.caching_hash_fn = caching_hash_fn
52 self.kv_cache_group_id = kv_cache_group_id
53 self._null_block = block_pool.null_block
54
@@ -130,14 +126,12 @@ class SingleTypeKVCacheManager(ABC):
126 req_blocks.extend(new_blocks)
127 return new_blocks
128
133 - def cache_blocks(self, request: Request, block_hashes: list[BlockHash],
134 - num_tokens: int) -> None:
129 + def cache_blocks(self, request: Request, num_tokens: int) -> None:
130 """
131 Cache the blocks for the request.
132
133 Args:
134 request: The request.
140 - block_hashes: The block hashes of the request.
135 num_tokens: The total number of tokens that need to be cached
136 (including tokens that are already cached).
137 """
@@ -147,12 +141,10 @@ class SingleTypeKVCacheManager(ABC):
141 self.block_pool.cache_full_blocks(
142 request=request,
143 blocks=self.req_to_blocks[request.request_id],
150 - block_hashes=block_hashes,
144 num_cached_blocks=num_cached_blocks,
145 num_full_blocks=num_full_blocks,
146 block_size=self.block_size,
147 kv_cache_group_id=self.kv_cache_group_id,
155 - hash_fn=self.caching_hash_fn,
148 )
149
150 self.num_cached_block[request.request_id] = num_full_blocks
vllm/v1/engine/core.py
+19 -3
@@ -25,9 +25,11 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
25 from vllm.tasks import POOLING_TASKS, SupportedTask
26 from vllm.transformers_utils.config import (
27 maybe_register_config_serialize_by_value)
28 -from vllm.utils import (decorate_logs, make_zmq_socket,
28 +from vllm.utils import (decorate_logs, get_hash_fn_by_name, make_zmq_socket,
29 resolve_obj_by_qualname, set_process_title)
30 -from vllm.v1.core.kv_cache_utils import (get_kv_cache_config,
30 +from vllm.v1.core.kv_cache_utils import (BlockHash, get_kv_cache_config,
31 + get_request_block_hasher,
32 + init_none_hash,
33 unify_kv_cache_configs)
34 from vllm.v1.core.sched.interface import SchedulerInterface
35 from vllm.v1.core.sched.output import SchedulerOutput
@@ -140,6 +142,19 @@ class EngineCore:
142 self.batch_queue_size)
143 self.batch_queue = queue.Queue(self.batch_queue_size)
144
145 + self.request_block_hasher: Optional[Callable[[Request],
146 + list[BlockHash]]] = None
147 + if (self.vllm_config.cache_config.enable_prefix_caching
148 + or self.scheduler.get_kv_connector() is not None):
149 +
150 + block_size = vllm_config.cache_config.block_size
151 + caching_hash_fn = get_hash_fn_by_name(
152 + vllm_config.cache_config.prefix_caching_hash_algo)
153 + init_none_hash(caching_hash_fn)
154 +
155 + self.request_block_hasher = get_request_block_hasher(
156 + block_size, caching_hash_fn)
157 +
158 def _initialize_kv_caches(
159 self, vllm_config: VllmConfig) -> tuple[int, int, KVCacheConfig]:
160 start = time.time()
@@ -417,7 +432,8 @@ class EngineCore:
432 request.mm_kwargs = self.mm_input_cache_server.get_and_update(
433 request.mm_kwargs, request.mm_hashes)
434
420 - req = Request.from_engine_core_request(request)
435 + req = Request.from_engine_core_request(request,
436 + self.request_block_hasher)
437 if req.use_structured_output:
438 # Note on thread safety: no race condition.
439 # `grammar_init` is only invoked in input processing thread. For
vllm/v1/request.py
+20 -2
@@ -3,7 +3,8 @@
3
4 import enum
5 import time
6 -from typing import TYPE_CHECKING, Any, Optional, Union
6 +from functools import partial
7 +from typing import TYPE_CHECKING, Any, Callable, Optional, Union
8
9 from vllm.multimodal.inputs import MultiModalKwargsItem, PlaceholderRange
10 from vllm.pooling_params import PoolingParams
@@ -16,6 +17,7 @@ from vllm.v1.utils import ConstantList
17
18 if TYPE_CHECKING:
19 from vllm.lora.request import LoRARequest
20 + from vllm.v1.core.kv_cache_utils import BlockHash
21
22
23 class Request:
@@ -36,6 +38,8 @@ class Request:
38 structured_output_request: Optional["StructuredOutputRequest"] = None,
39 cache_salt: Optional[str] = None,
40 priority: int = 0,
41 + block_hasher: Optional[Callable[["Request"],
42 + list["BlockHash"]]] = None,
43 ) -> None:
44 self.request_id = request_id
45 self.client_index = client_index
@@ -108,8 +112,18 @@ class Request:
112 # indicates that the output is corrupted
113 self.num_nans_in_logits = 0
114
115 + self.block_hashes: list[BlockHash] = []
116 + self.get_hash_new_full_blocks: Optional[Callable[
117 + [], list[BlockHash]]] = None
118 + if block_hasher is not None:
119 + self.get_hash_new_full_blocks = partial(block_hasher, self)
120 + self.block_hashes = self.get_hash_new_full_blocks()
121 +
122 @classmethod
112 - def from_engine_core_request(cls, request: EngineCoreRequest) -> "Request":
123 + def from_engine_core_request(
124 + cls, request: EngineCoreRequest,
125 + block_hasher: Optional[Callable[["Request"], list["BlockHash"]]]
126 + ) -> "Request":
127 if request.mm_kwargs is not None:
128 assert is_list_of(request.mm_kwargs, MultiModalKwargsItem), (
129 "mm_kwargs was not updated in EngineCore.add_request")
@@ -131,6 +145,7 @@ class Request:
145 if request.sampling_params else None,
146 cache_salt=request.cache_salt,
147 priority=request.priority,
148 + block_hasher=block_hasher,
149 )
150
151 def append_output_token_ids(
@@ -144,6 +159,9 @@ class Request:
159 self._output_token_ids.extend(token_ids)
160 self._all_token_ids.extend(token_ids)
161
162 + if self.get_hash_new_full_blocks is not None:
163 + self.block_hashes.extend(self.get_hash_new_full_blocks())
164 +
165 @property
166 def is_output_corrupted(self) -> bool:
167 return self.num_nans_in_logits > 0