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
[Core] Use sha256 bytes instead of BlockHash to reduce GC overhead (#23673)
Signed-off-by: linzebing <linzebing1995@gmail.com>
Zebing Lin committed
Sep 9, 2025 at 00:34 UTC
82dfb12e52ded608da2a458b2da5f48ba1acecdf
15 files changed
+298
-283
examples/online_serving/kv_events_subscriber.py
+5
-3
@@ -6,6 +6,8 @@ import msgspec
6
import zmq
7
from msgspec.msgpack import Decoder
8
9
+from vllm.v1.core.kv_cache_utils import BlockHash
10
+
11
12
#
13
# Types copied from vllm.distributed.kv_events
@@ -22,8 +24,8 @@ class KVCacheEvent(
24
25
26
class BlockStored(KVCacheEvent):
25
- block_hashes: list[int]
26
- parent_block_hash: Optional[int]
27
+ block_hashes: list[BlockHash]
28
+ parent_block_hash: Optional[BlockHash]
29
token_ids: list[int]
30
block_size: int
31
lora_id: Optional[int]
@@ -31,7 +33,7 @@ class BlockStored(KVCacheEvent):
33
34
35
class BlockRemoved(KVCacheEvent):
34
- block_hashes: list[int]
36
+ block_hashes: list[BlockHash]
37
medium: Optional[str]
38
39
tests/utils_/test_utils.py
+9
-11
@@ -835,22 +835,20 @@ def test_model_specification(parser_with_config, cli_config_file,
835
836
@pytest.mark.parametrize("input", [(), ("abc", ), (None, ),
837
(None, bool, [1, 2, 3])])
838
-@pytest.mark.parametrize("output", [0, 1, 2])
839
-def test_sha256(input: tuple, output: int):
840
- hash = sha256(input)
841
- assert hash is not None
842
- assert isinstance(hash, int)
843
- assert hash != 0
838
+def test_sha256(input: tuple):
839
+ digest = sha256(input)
840
+ assert digest is not None
841
+ assert isinstance(digest, bytes)
842
+ assert digest != b""
843
845
- bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
846
- assert hash == int.from_bytes(hashlib.sha256(bytes).digest(),
847
- byteorder="big")
844
+ input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
845
+ assert digest == hashlib.sha256(input_bytes).digest()
846
847
# hashing again, returns the same value
850
- assert hash == sha256(input)
848
+ assert digest == sha256(input)
849
850
# hashing different input, returns different value
853
- assert hash != sha256(input + (1, ))
851
+ assert digest != sha256(input + (1, ))
852
853
854
@pytest.mark.parametrize(
tests/v1/core/test_kv_cache_utils.py
+27
-38
@@ -6,20 +6,22 @@ from typing import Callable, Optional
6
import pytest
7
import torch
8
9
+import vllm.v1.core.kv_cache_utils as kv_cache_utils
10
from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
11
from vllm.multimodal.inputs import (MultiModalFeatureSpec,
12
MultiModalKwargsItem, PlaceholderRange)
13
from vllm.sampling_params import SamplingParams
13
-from vllm.utils import GiB_bytes, sha256, sha256_cbor_64bit
14
+from vllm.utils import GiB_bytes, sha256, sha256_cbor
15
from vllm.v1.core.kv_cache_manager import KVCacheManager
16
# disable yapf here as it formats differently than isort such that both fail
17
# yapf: disable
18
from vllm.v1.core.kv_cache_utils import (
18
- FreeKVCacheBlockQueue, KVCacheBlock, PrefixCachingMetrics,
19
+ BlockHash, 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
get_request_block_hasher, hash_block_tokens, init_none_hash,
22
- is_kv_cache_type_uniform, unify_kv_cache_configs)
23
+ is_kv_cache_type_uniform, make_block_hash_with_group_id,
24
+ unify_kv_cache_configs)
25
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
26
KVCacheGroupSpec, KVCacheTensor,
27
SlidingWindowSpec)
@@ -88,7 +90,7 @@ def new_sliding_window_spec(block_size=16,
90
sliding_window=sliding_window)
91
92
91
-@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
93
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
94
def test_none_hash(monkeypatch, hash_fn):
95
import vllm.v1.core.kv_cache_utils
96
@@ -98,8 +100,8 @@ def test_none_hash(monkeypatch, hash_fn):
100
reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
101
reloaded_kv_cache_utils.init_none_hash(hash_fn)
102
assert reloaded_kv_cache_utils.NONE_HASH is not None
101
- assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
102
- assert reloaded_kv_cache_utils.NONE_HASH != 0
103
+ assert isinstance(reloaded_kv_cache_utils.NONE_HASH, bytes)
104
+ assert reloaded_kv_cache_utils.NONE_HASH != b""
105
106
# case 2: PYTHONHASHSEED is set, use the seed and hash_fn
107
with monkeypatch.context() as m:
@@ -107,12 +109,11 @@ def test_none_hash(monkeypatch, hash_fn):
109
reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
110
reloaded_kv_cache_utils.init_none_hash(hash_fn)
111
assert reloaded_kv_cache_utils.NONE_HASH is not None
110
- assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
112
+ assert isinstance(reloaded_kv_cache_utils.NONE_HASH, bytes)
113
assert hash_fn('python hash seed') == reloaded_kv_cache_utils.NONE_HASH
114
115
116
def test_kv_cache_block():
115
- import vllm.v1.core.kv_cache_utils
117
118
# Test KVCacheBlock initialization
119
block = KVCacheBlock(block_id=0)
@@ -127,8 +128,7 @@ def test_kv_cache_block():
128
assert block.ref_cnt == 0
129
130
# Test block hash setting and resetting
130
- block_hash = vllm.v1.core.kv_cache_utils.BlockHash(hash_value=123,
131
- token_ids=(1, 2, 3))
131
+ block_hash = make_block_hash_with_group_id(BlockHash(b"abc"), 0)
132
block.block_hash = block_hash
133
assert block.block_hash == block_hash
134
@@ -407,27 +407,23 @@ def test_generate_block_hash_extra_keys_cache_salt():
407
assert next_mm_idx == 1
408
409
410
-@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
410
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
411
def test_hash_block_tokens(hash_fn):
412
- import vllm.v1.core.kv_cache_utils
412
init_none_hash(hash_fn)
414
- parent_block_hash = 123
413
+ parent_block_hash = BlockHash(b"123")
414
curr_block_token_ids = (1, 2, 3)
415
extra_keys = ("key1", "key2")
416
417
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
418
curr_block_token_ids, extra_keys)
420
- assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHash)
421
- assert block_hash.hash_value == hash_fn(
422
- (parent_block_hash, curr_block_token_ids, extra_keys))
423
- assert block_hash.token_ids == curr_block_token_ids
424
- assert block_hash.extra_keys == extra_keys
419
+ expected = hash_fn((parent_block_hash, curr_block_token_ids, extra_keys))
420
+ assert block_hash == expected
421
422
427
-@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
423
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
424
def test_request_block_hasher(hash_fn):
429
- import vllm.v1.core.kv_cache_utils
430
- init_none_hash(hash_fn)
425
+ kv_cache_utils.init_none_hash(hash_fn)
426
+
427
request = make_request(
428
request_id="0",
429
prompt_token_ids=[_ for _ in range(6)],
@@ -442,19 +438,13 @@ def test_request_block_hasher(hash_fn):
438
439
block_hashes = request.block_hashes
440
assert len(block_hashes) == 2
445
- assert isinstance(block_hashes[0], vllm.v1.core.kv_cache_utils.BlockHash)
446
- assert isinstance(block_hashes[1], vllm.v1.core.kv_cache_utils.BlockHash)
447
-
448
- # Check the first block
449
- assert block_hashes[0].token_ids == (0, 1, 2)
450
- assert block_hashes[0].extra_keys == ("hash1", )
441
+ assert block_hashes[0] == hash_fn(
442
+ (kv_cache_utils.NONE_HASH, (0, 1, 2), ("hash1", )))
443
+ assert block_hashes[1] == hash_fn(
444
+ (block_hashes[0], (3, 4, 5), ("hash2", )))
445
452
- # Check the second block
453
- assert block_hashes[1].token_ids == (3, 4, 5)
454
- assert block_hashes[1].extra_keys == ("hash2", )
446
456
-
457
-@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
447
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
448
def test_hash_tokens_different_mm_input(hash_fn):
449
init_none_hash(hash_fn)
450
@@ -484,9 +474,9 @@ def test_hash_tokens_different_mm_input(hash_fn):
474
assert block_hashes1[1] != block_hashes2[1]
475
476
487
-@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
477
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
478
def test_hash_request_tokens_no_mm_inputs(hash_fn):
489
- init_none_hash(hash_fn)
479
+ kv_cache_utils.init_none_hash(hash_fn)
480
481
request = make_request(
482
request_id="0",
@@ -500,10 +490,9 @@ def test_hash_request_tokens_no_mm_inputs(hash_fn):
490
block_hashes = request.block_hashes
491
492
assert len(block_hashes) == 2
503
- assert block_hashes[0].token_ids == (0, 1, 2)
504
- assert block_hashes[0].extra_keys is None
505
- assert block_hashes[1].token_ids == (3, 4, 5)
506
- assert block_hashes[1].extra_keys is None
493
+ assert block_hashes[0] == hash_fn(
494
+ (kv_cache_utils.NONE_HASH, (0, 1, 2), None))
495
+ assert block_hashes[1] == hash_fn((block_hashes[0], (3, 4, 5), None))
496
497
498
def test_metrics():
tests/v1/core/test_prefix_caching.py
+122
-103
@@ -8,17 +8,19 @@ from typing import Callable, Optional
8
import pytest
9
import torch
10
11
+import vllm.v1.core.kv_cache_utils as kv_cache_utils
12
from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved
13
from vllm.multimodal.inputs import (MultiModalFeatureSpec,
14
MultiModalKwargsItem, PlaceholderRange)
15
from vllm.sampling_params import SamplingParams
15
-from vllm.utils import sha256, sha256_cbor_64bit
16
+from vllm.utils import sha256, sha256_cbor
17
from vllm.v1.core.block_pool import BlockPool
18
from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
18
-from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
19
- KVCacheBlock,
19
+from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
20
+ get_block_hash, get_group_id,
21
get_request_block_hasher,
21
- hash_block_tokens, init_none_hash)
22
+ hash_block_tokens, init_none_hash,
23
+ make_block_hash_with_group_id)
24
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
25
KVCacheGroupSpec, SlidingWindowSpec)
26
@@ -101,8 +103,10 @@ def make_kv_cache_config_hybrid_model(block_size: int,
103
)
104
105
104
-@pytest.mark.parametrize("hash_algo", ["sha256", "sha256_cbor_64bit", "hash"])
105
-def test_prefill(hash_algo):
106
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
107
+def test_prefill(hash_fn):
108
+ init_none_hash(hash_fn)
109
+
110
block_size = 16
111
manager = KVCacheManager(
112
make_kv_cache_config(block_size, 11),
@@ -110,10 +114,6 @@ def test_prefill(hash_algo):
114
enable_caching=True,
115
)
116
113
- # choose the hash function according to the parameter
114
- hash_fn = (sha256_cbor_64bit if hash_algo == "sha256_cbor_64bit" else
115
- sha256 if hash_algo == "sha256" else hash)
116
-
117
# Complete 3 blocks (48 tokens)
118
common_token_ids = [i for i in range(3) for _ in range(16)]
119
@@ -137,10 +137,12 @@ def test_prefill(hash_algo):
137
block_tokens = tuple(all_token_ids[(block_id - 1) * 16:block_id * 16])
138
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
139
block_tokens)
140
- assert manager.block_pool.blocks[
141
- block_id].block_hash.block_hash == block_hash
140
+ blk_hash = manager.block_pool.blocks[block_id].block_hash
141
+ assert blk_hash is not None
142
+ assert get_block_hash(blk_hash) == block_hash
143
+ assert get_group_id(blk_hash) == 0
144
assert manager.block_pool.blocks[block_id].ref_cnt == 1
143
- parent_block_hash = block_hash.hash_value
145
+ parent_block_hash = block_hash
146
147
# Check partial block metadata
148
for block_id in (4, ):
@@ -233,7 +235,7 @@ def test_prefill_hybrid_model():
235
enable_caching=True,
236
)
237
236
- hash_fn = hash
238
+ hash_fn = sha256
239
240
# Complete 3 blocks (48 tokens)
241
common_token_ids = [i for i in range(3) for _ in range(block_size)]
@@ -260,11 +262,13 @@ def test_prefill_hybrid_model():
262
block_tokens = tuple(all_token_ids[(length - 1) * 16:length * 16])
263
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
264
block_tokens)
263
- for block_id in block_ids:
264
- assert manager.block_pool.blocks[
265
- block_id].block_hash.block_hash == block_hash
265
+ for group_id, block_id in enumerate(block_ids):
266
+ blk_hash = manager.block_pool.blocks[block_id].block_hash
267
+ assert blk_hash is not None
268
+ assert get_block_hash(blk_hash) == block_hash
269
+ assert get_group_id(blk_hash) == group_id
270
assert manager.block_pool.blocks[block_id].ref_cnt == 1
267
- parent_block_hash = block_hash.hash_value
271
+ parent_block_hash = block_hash
272
273
# Check partial block metadata
274
for block_id in (4, 8, 12):
@@ -298,11 +302,10 @@ def test_prefill_hybrid_model():
302
cached_block_hash_to_block_bak = copy.copy(
303
manager.block_pool.cached_block_hash_to_block)
304
301
- def test_partial_request_hit(request_id: str,
302
- hash_to_evict: list[BlockHashWithGroupId],
305
+ def test_partial_request_hit(request_id: str, hash_to_evict: list[bytes],
306
expect_hit_length: int):
307
req = make_request(request_id, common_token_ids + unique_token_ids,
305
- block_size, hash)
308
+ block_size, sha256)
309
for hash_with_group_id in hash_to_evict:
310
manager.block_pool.cached_block_hash_to_block.pop(
311
hash_with_group_id)
@@ -319,33 +322,32 @@ def test_prefill_hybrid_model():
322
323
# Evict the blocks outside sliding window, does not affect the hit length.
324
test_partial_request_hit("2", [
322
- BlockHashWithGroupId(block_hashes[0], 1),
323
- BlockHashWithGroupId(block_hashes[0], 2)
325
+ make_block_hash_with_group_id(block_hashes[0], 1),
326
+ make_block_hash_with_group_id(block_hashes[0], 2)
327
], 3)
328
329
# Evict the first block of full attention, makes total cache miss.
327
- test_partial_request_hit("3", [
328
- BlockHashWithGroupId(block_hashes[0], 0),
329
- ], 0)
330
+ test_partial_request_hit(
331
+ "3", [make_block_hash_with_group_id(block_hashes[0], 0)], 0)
332
333
# Evict the last block of all layers, reduces the hit length to 2.
334
test_partial_request_hit("4", [
333
- BlockHashWithGroupId(block_hashes[2], 0),
334
- BlockHashWithGroupId(block_hashes[2], 1),
335
- BlockHashWithGroupId(block_hashes[2], 2),
335
+ make_block_hash_with_group_id(block_hashes[2], 0),
336
+ make_block_hash_with_group_id(block_hashes[2], 1),
337
+ make_block_hash_with_group_id(block_hashes[2], 2),
338
], 2)
339
340
# Evict the last block of full attention, reduces the hit length to 2.
339
- test_partial_request_hit("5", [BlockHashWithGroupId(block_hashes[2], 0)],
340
- 2)
341
+ test_partial_request_hit(
342
+ "5", [make_block_hash_with_group_id(block_hashes[2], 0)], 2)
343
344
# Evict the last block of sliding window, reduces the hit length to 2.
343
- test_partial_request_hit("6", [BlockHashWithGroupId(block_hashes[2], 1)],
344
- 2)
345
+ test_partial_request_hit(
346
+ "6", [make_block_hash_with_group_id(block_hashes[2], 1)], 2)
347
348
# Evict the last block of sliding window, reduces the hit length to 2.
347
- test_partial_request_hit("7", [BlockHashWithGroupId(block_hashes[2], 2)],
348
- 2)
349
+ test_partial_request_hit(
350
+ "7", [make_block_hash_with_group_id(block_hashes[2], 2)], 2)
351
352
# Evict different set of blocks for full attention and sliding window makes
353
# total cache miss.
@@ -353,9 +355,9 @@ def test_prefill_hybrid_model():
355
# The cache hit length of sliding window is 2 * block_size.
356
# Then it is cache miss as the two type of layers have different hit length.
357
test_partial_request_hit("8", [
356
- BlockHashWithGroupId(block_hashes[2], 0),
357
- BlockHashWithGroupId(block_hashes[0], 1),
358
- BlockHashWithGroupId(block_hashes[0], 2),
358
+ make_block_hash_with_group_id(block_hashes[2], 0),
359
+ make_block_hash_with_group_id(block_hashes[0], 1),
360
+ make_block_hash_with_group_id(block_hashes[0], 2),
361
], 0)
362
363
@@ -372,8 +374,8 @@ def test_prefill_plp():
374
max_model_len=8192,
375
enable_caching=True,
376
)
375
- # the default hash function is hash
376
- hash_fn = hash
377
+ # the default hash function is sha256
378
+ hash_fn = sha256
379
380
# Complete 3 blocks (48 tokens)
381
common_token_ids = [i for i in range(3) for _ in range(16)]
@@ -404,10 +406,12 @@ def test_prefill_plp():
406
block_tokens = tuple(all_token_ids[(block_id - 1) * 16:block_id * 16])
407
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
408
block_tokens)
407
- assert manager.block_pool.blocks[
408
- block_id].block_hash.block_hash == block_hash
409
+ blk_hash = (manager.block_pool.blocks[block_id].block_hash)
410
+ assert blk_hash is not None
411
+ assert get_block_hash(blk_hash) == block_hash
412
+ assert get_group_id(blk_hash) == 0
413
assert manager.block_pool.blocks[block_id].ref_cnt == 1
410
- parent_block_hash = block_hash.hash_value
414
+ parent_block_hash = block_hash
415
416
# Check partial block metadata
417
for block_id in (4, ):
@@ -493,7 +497,7 @@ def test_decode():
497
# Incomplete 1 block (7 tokens)
498
unique_token_ids = [3] * 7
499
req0 = make_request("0", common_token_ids + unique_token_ids, block_size,
496
- hash)
500
+ sha256)
501
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
502
assert not computed_blocks.blocks[0]
503
assert num_computed_tokens == 0
@@ -538,7 +542,7 @@ def test_evict():
542
)
543
544
last_token_id = 5 * 16 + 7
541
- req0 = make_request("0", list(range(last_token_id)), block_size, hash)
545
+ req0 = make_request("0", list(range(last_token_id)), block_size, sha256)
546
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
547
assert not computed_blocks.blocks[0]
548
assert num_computed_tokens == 0
@@ -550,7 +554,7 @@ def test_evict():
554
# 3 blocks.
555
req1 = make_request("1", list(range(last_token_id,
556
last_token_id + 3 * 16)), block_size,
553
- hash)
557
+ sha256)
558
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
559
assert not computed_blocks.blocks[0]
560
assert num_computed_tokens == 0
@@ -572,7 +576,7 @@ def test_evict():
576
] == [10, 6, 5, 4, 3, 2, 1, 9, 8, 7]
577
578
# Touch the first 2 blocks.
575
- req2 = make_request("2", list(range(2 * 16 + 3)), block_size, hash)
579
+ req2 = make_request("2", list(range(2 * 16 + 3)), block_size, sha256)
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
@@ -597,7 +601,7 @@ def test_hash_block_correct_reuse():
601
602
# Allocate 1 block and cache it.
603
num_tokens = block_size * 1
600
- req = make_request("0", list(range(num_tokens)), block_size, hash)
604
+ req = make_request("0", list(range(num_tokens)), block_size, sha256)
605
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
606
assert not computed_blocks.blocks[0]
607
assert num_computed_tokens == 0
@@ -611,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.
614
- req = make_request("1", list(range(num_tokens - 1)), block_size, hash)
618
+ req = make_request("1", list(range(num_tokens - 1)), block_size, sha256)
619
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
620
assert not computed_blocks.blocks[0]
621
assert num_computed_tokens == 0
@@ -638,7 +642,7 @@ def test_computed_blocks_not_evicted():
642
643
# Allocate a block and cache it.
644
num_tokens = block_size * 1
641
- req0 = make_request("0", list(range(num_tokens)), block_size, hash)
645
+ req0 = make_request("0", list(range(num_tokens)), block_size, sha256)
646
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
647
assert not computed_blocks.blocks[0]
648
assert num_computed_tokens == 0
@@ -650,7 +654,7 @@ def test_computed_blocks_not_evicted():
654
655
# Allocate another block.
656
req1 = make_request("1", list(range(num_tokens, num_tokens * 2)),
653
- block_size, hash)
657
+ block_size, sha256)
658
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
659
assert not computed_blocks.blocks[0]
660
assert num_computed_tokens == 0
@@ -666,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.
669
- req2 = make_request("2", list(range(num_tokens * 2)), block_size, hash)
673
+ req2 = make_request("2", list(range(num_tokens * 2)), block_size, sha256)
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
@@ -691,7 +695,7 @@ def test_basic_prefix_caching_disabled():
695
)
696
697
req1 = make_request("1", list(range(10)), block_size,
694
- hash) # 2 blocks and some more
698
+ sha256) # 2 blocks and some more
699
700
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
701
assert not computed_blocks.blocks[0]
@@ -706,7 +710,7 @@ def test_basic_prefix_caching_disabled():
710
711
# No caching.
712
req2 = make_request("2", list(range(16)), block_size,
709
- hash) # shared prefix
713
+ sha256) # 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
@@ -716,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.
719
- req3 = make_request("3", list(range(4)), block_size, hash)
723
+ req3 = make_request("3", list(range(4)), block_size, sha256)
724
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
725
assert not computed_blocks.blocks[0]
726
assert num_computed_tokens == 0
@@ -726,7 +730,7 @@ def test_basic_prefix_caching_disabled():
730
assert not blocks
731
732
729
-@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor_64bit, hash])
733
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
734
def test_cache_blocks(hash_fn):
735
"""
736
This is a unit test that tests the correctness of the _cache_full_blocks
@@ -787,7 +791,7 @@ 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]
790
- req = make_request("0", list(range(14)), block_size, hash)
794
+ req = make_request("0", list(range(14)), block_size, sha256)
795
796
# Cache the blocks for group 0.
797
blocks = [KVCacheBlock(block_id=i) for i in range(2)]
@@ -845,6 +849,8 @@ def test_mm_prefix_caching():
849
"""
850
This tests that the multi-modal prefix caching is correct.
851
"""
852
+ kv_cache_utils.init_none_hash(sha256)
853
+
854
block_size = 16
855
manager = KVCacheManager(
856
make_kv_cache_config(block_size, 11),
@@ -874,23 +880,30 @@ def test_mm_prefix_caching():
880
req0 = make_request("0",
881
all_token_ids,
882
block_size,
877
- hash,
883
+ sha256,
884
mm_positions=mm_positions,
885
mm_hashes=mm_hashes)
886
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
887
882
- # Completed block should have hashes with extra keys.
888
+ # Completed block should have hashes
889
assert not computed_blocks.blocks[0]
890
assert num_computed_tokens == 0
891
block_hashes = req0.block_hashes
892
assert len(block_hashes) == 3
887
- assert block_hashes[0].extra_keys == ("aaa", )
888
- assert block_hashes[1].extra_keys == ("aaa", "bbb")
889
- assert block_hashes[2].extra_keys == ("bbb", )
893
+ assert block_hashes[0] == sha256(
894
+ (kv_cache_utils.NONE_HASH, tuple(all_token_ids[:block_size]),
895
+ ("aaa", )))
896
+ assert block_hashes[1] == sha256(
897
+ (block_hashes[0], tuple(all_token_ids[block_size:block_size * 2]),
898
+ ("aaa", "bbb")))
899
+ assert block_hashes[2] == sha256(
900
+ (block_hashes[1], tuple(all_token_ids[block_size * 2:block_size * 3]),
901
+ ("bbb", )))
902
903
blocks = manager.allocate_slots(req0, 59,
904
len(computed_blocks.blocks[0]) * 16,
905
computed_blocks)
906
+ assert blocks is not None
907
assert blocks.get_block_ids() == ([1, 2, 3, 4], )
908
req0.num_computed_tokens = 59
909
@@ -901,10 +914,10 @@ def test_mm_prefix_caching():
914
len(computed_blocks.blocks[0]) * 16,
915
computed_blocks)
916
assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
904
-
905
- # The just completed block should have hashes with extra keys.
917
assert len(block_hashes) == 4
907
- assert block_hashes[3].extra_keys == ("ccc", )
918
+ assert block_hashes[3] == sha256(
919
+ (block_hashes[2], tuple(all_token_ids[3 * block_size:] + [8] * 5),
920
+ ("ccc", )))
921
922
# Cache hit.
923
unique_token_ids = [-1] * 7 + [200] * 5
@@ -916,7 +929,7 @@ def test_mm_prefix_caching():
929
req1 = make_request("1",
930
all_token_ids,
931
block_size,
919
- hash,
932
+ sha256,
933
mm_positions=mm_positions,
934
mm_hashes=mm_hashes)
935
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
@@ -929,6 +942,8 @@ def test_cache_key_salting():
942
This tests that cache salts are applied during hashing and the cache
943
is separated cache as expected.
944
"""
945
+ kv_cache_utils.init_none_hash(sha256)
946
+
947
block_size = 16
948
manager = KVCacheManager(
949
make_kv_cache_config(block_size, 11),
@@ -939,21 +954,26 @@ def test_cache_key_salting():
954
# 3 complete blocks and an incomplete block with 11 tokens.
955
common_token_ids = [i for i in range(3) for _ in range(block_size)]
956
token_ids = common_token_ids + [3] * 11
942
- req0 = make_request("0", token_ids, block_size, hash, cache_salt="salt1")
957
+ req0 = make_request("0", token_ids, block_size, sha256, cache_salt="salt1")
958
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
959
945
- # Completed block should have hashes with extra keys.
960
+ # Completed block should have hashes
961
assert not computed_blocks.blocks[0]
962
assert num_computed_tokens == 0
963
block_hashes = req0.block_hashes
964
assert len(block_hashes) == 3
950
- assert block_hashes[0].extra_keys == ("salt1", )
951
- assert block_hashes[1].extra_keys is None
952
- assert block_hashes[2].extra_keys is None
965
+ assert block_hashes[0] == sha256(
966
+ (kv_cache_utils.NONE_HASH, tuple(token_ids[:block_size]), ("salt1", )))
967
+ assert block_hashes[1] == sha256(
968
+ (block_hashes[0], tuple(token_ids[block_size:block_size * 2]), None))
969
+ assert block_hashes[2] == sha256(
970
+ (block_hashes[1], tuple(token_ids[block_size * 2:block_size * 3]),
971
+ None))
972
973
blocks = manager.allocate_slots(req0, 59,
974
len(computed_blocks.blocks[0]) * 16,
975
computed_blocks)
976
+ assert blocks is not None
977
assert blocks.get_block_ids() == ([1, 2, 3, 4], )
978
req0.num_computed_tokens = 59
979
@@ -964,14 +984,13 @@ def test_cache_key_salting():
984
len(computed_blocks.blocks[0]) * 16,
985
computed_blocks)
986
assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
967
-
968
- # Now one more block that should not have extra keys.
987
assert len(block_hashes) == 4
970
- assert block_hashes[3].extra_keys is None
988
+ assert block_hashes[3] == sha256(
989
+ (block_hashes[2], tuple(token_ids[3 * block_size:] + [8] * 5), None))
990
991
# Test cache hit with a new request that has the same salt.
992
token_ids = common_token_ids + [4] * 11
974
- req1 = make_request("1", token_ids, block_size, hash, cache_salt="salt1")
993
+ req1 = make_request("1", token_ids, block_size, sha256, cache_salt="salt1")
994
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
995
# Should match only a prefix of 3 blocks.
996
assert len(computed_blocks.blocks[0]) == 3
@@ -979,13 +998,19 @@ def test_cache_key_salting():
998
999
# Test cache miss with same content but different salt.
1000
token_ids = common_token_ids + [4] * 11
982
- req2 = make_request("2", token_ids, block_size, hash, cache_salt="salt2")
1001
+ req2 = make_request("2", token_ids, block_size, sha256, cache_salt="salt2")
1002
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
1003
assert len(computed_blocks.blocks[0]) == 0
1004
assert num_computed_tokens == 0
1005
block_hashes = req2.block_hashes
1006
assert len(block_hashes) == 3
988
- assert block_hashes[0].extra_keys == ("salt2", )
1007
+ assert block_hashes[0] == sha256(
1008
+ (kv_cache_utils.NONE_HASH, tuple(token_ids[:block_size]), ("salt2", )))
1009
+ assert block_hashes[1] == sha256(
1010
+ (block_hashes[0], tuple(token_ids[block_size:block_size * 2]), None))
1011
+ assert block_hashes[2] == sha256(
1012
+ (block_hashes[1], tuple(token_ids[block_size * 2:block_size * 3]),
1013
+ None))
1014
1015
1016
def test_prefill_not_enough_free_blocks_with_computed_blocks():
@@ -1004,7 +1029,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1029
# Complete 3 blocks (48 tokens)
1030
# | Common-0 | Common-1 | Common-2 | ... |
1031
common_token_ids = [i for i in range(3) for _ in range(16)]
1007
- req0 = make_request("0", common_token_ids, block_size, hash)
1032
+ req0 = make_request("0", common_token_ids, block_size, sha256)
1033
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
1034
assert not computed_blocks.blocks[0]
1035
assert num_computed_tokens == 0
@@ -1015,7 +1040,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1040
req0.request_id]
1041
1042
# | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |
1018
- req1 = make_request("1", common_token_ids * 2, block_size, hash)
1043
+ req1 = make_request("1", common_token_ids * 2, block_size, sha256)
1044
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
1045
assert computed_blocks.blocks[0] == block_part0
1046
assert num_computed_tokens == 3 * 16
@@ -1032,7 +1057,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1057
1058
# | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
1059
# | Req1-5(F)| Req2-0 | Req2-1 | ... |
1035
- req2 = make_request("2", [7] * block_size * 2, block_size, hash)
1060
+ req2 = make_request("2", [7] * block_size * 2, block_size, sha256)
1061
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
1062
assert not computed_blocks.blocks[0]
1063
assert num_computed_tokens == 0
@@ -1044,7 +1069,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1069
# but it cannot be allocated due to insufficient free blocks (2).
1070
# In this case, the ref_cnt of the computed blocks should not be changed.
1071
assert manager.block_pool.free_block_queue.num_free_blocks == 5
1047
- req3 = make_request("3", common_token_ids * 3, block_size, hash)
1072
+ req3 = make_request("3", common_token_ids * 3, block_size, sha256)
1073
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
1074
assert computed_blocks.blocks[0] == block_part1
1075
assert num_computed_tokens == 6 * 16
@@ -1069,13 +1094,13 @@ def test_reset_prefix_cache():
1094
full_block_token_ids = [i for i in range(3) for _ in range(16)]
1095
unique_token_ids = [3] * 7
1096
all_token_ids = full_block_token_ids + unique_token_ids
1072
- req0 = make_request("0", all_token_ids, block_size, hash)
1097
+ req0 = make_request("0", all_token_ids, block_size, sha256)
1098
blocks = manager.allocate_slots(req0, 55)
1099
assert blocks.get_block_ids() == ([1, 2, 3, 4], )
1100
1101
unique_token_ids = [4] * 7
1102
all_token_ids = full_block_token_ids + unique_token_ids
1078
- req1 = make_request("1", all_token_ids, block_size, hash)
1103
+ req1 = make_request("1", all_token_ids, block_size, sha256)
1104
computed_blocks, _ = manager.get_computed_blocks(req1)
1105
assert len(req1.block_hashes) == 3
1106
assert len(computed_blocks.blocks[0]) == 3
@@ -1109,7 +1134,7 @@ def test_prefix_cache_stats_disabled():
1134
assert manager.prefix_cache_stats is None
1135
1136
# Call all functions that check whether log_stats is disabled.
1112
- req = make_request("0", list(range(16)), block_size, hash)
1137
+ req = make_request("0", list(range(16)), block_size, sha256)
1138
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
1139
assert not computed_blocks.blocks[0]
1140
assert num_computed_tokens == 0
@@ -1124,15 +1149,9 @@ def test_prefix_cache_stats_disabled():
1149
1150
def test_maybe_evict_cached_block():
1151
pool = BlockPool(num_gpu_blocks=4, enable_caching=True)
1127
- block_hash0 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=10,
1128
- token_ids=(100, )),
1129
- group_id=1000)
1130
- block_hash1 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=20,
1131
- token_ids=(200, )),
1132
- group_id=2000)
1133
- block_hash2 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=30,
1134
- token_ids=(300, )),
1135
- group_id=3000)
1152
+ block_hash0 = make_block_hash_with_group_id(BlockHash(b"10"), 1000)
1153
+ block_hash1 = make_block_hash_with_group_id(BlockHash(b"20"), 2000)
1154
+ block_hash2 = make_block_hash_with_group_id(BlockHash(b"30"), 3000)
1155
block_hashes = [
1156
block_hash0,
1157
block_hash1,
@@ -1206,7 +1225,7 @@ def test_kv_cache_events(blocks_to_cache: int):
1225
)
1226
1227
num_tokens = block_size * blocks_to_cache
1209
- req0 = make_request("0", list(range(num_tokens)), block_size, hash)
1228
+ req0 = make_request("0", list(range(num_tokens)), block_size, sha256)
1229
_ = manager.allocate_slots(req0, num_tokens)
1230
events = manager.take_events()
1231
@@ -1222,7 +1241,7 @@ def test_kv_cache_events(blocks_to_cache: int):
1241
# Should see block_to_cache number of removed block events and a new block
1242
# stored event
1243
manager.free(req0)
1225
- req1 = make_request("1", list(range(num_tokens)), block_size, hash)
1244
+ req1 = make_request("1", list(range(num_tokens)), block_size, sha256)
1245
_ = manager.allocate_slots(req1, num_tokens)
1246
events = manager.take_events()
1247
@@ -1256,7 +1275,7 @@ def test_eagle_enabled_removes_last_block():
1275
1276
# Request with 3 full blocks (48 tokens)
1277
token_ids = [0] * (3 * block_size)
1259
- req = make_request("divisible_request", token_ids, block_size, hash)
1278
+ req = make_request("divisible_request", token_ids, block_size, sha256)
1279
1280
# Prime the cache
1281
computed_blocks, _ = manager.get_computed_blocks(req)
@@ -1266,7 +1285,7 @@ def test_eagle_enabled_removes_last_block():
1285
manager.free(req)
1286
1287
# New request with same tokens + Eagle enabled
1269
- req_eagle = make_request("eagle_divisible", token_ids, block_size, hash)
1288
+ req_eagle = make_request("eagle_divisible", token_ids, block_size, sha256)
1289
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1290
1291
# Should retain 1 block:
@@ -1287,7 +1306,7 @@ def test_eagle_with_partial_blocks():
1306
)
1307
# 2 full blocks + 5 tokens (non-divisible length)
1308
token_ids = [0] * (2 * block_size + 5)
1290
- req = make_request("partial_block_test", token_ids, block_size, hash)
1309
+ req = make_request("partial_block_test", token_ids, block_size, sha256)
1310
1311
# Prime the cache
1312
computed_blocks, _ = manager.get_computed_blocks(req)
@@ -1297,7 +1316,7 @@ def test_eagle_with_partial_blocks():
1316
manager.free(req)
1317
1318
# New request with Eagle enabled
1300
- req_eagle = make_request("partial_eagle", token_ids, block_size, hash)
1319
+ req_eagle = make_request("partial_eagle", token_ids, block_size, sha256)
1320
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1321
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
1322
assert len(computed_blocks.blocks[0]) == 1
@@ -1328,7 +1347,7 @@ def test_eagle_with_sliding_window():
1347
1348
# 2 full blocks + 5 tokens (non-divisible length)
1349
token_ids = [0] * (2 * block_size + 5)
1331
- req = make_request("partial_block_test", token_ids, block_size, hash)
1350
+ req = make_request("partial_block_test", token_ids, block_size, sha256)
1351
1352
# Prime the cache
1353
computed_blocks, _ = manager.get_computed_blocks(req)
@@ -1341,7 +1360,7 @@ def test_eagle_with_sliding_window():
1360
manager.free(req)
1361
1362
# New request with Eagle enabled
1344
- req_eagle = make_request("partial_eagle", token_ids, block_size, hash)
1363
+ req_eagle = make_request("partial_eagle", token_ids, block_size, sha256)
1364
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1365
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
1366
assert len(computed_blocks.blocks[0]) == 1
@@ -1351,11 +1370,11 @@ def test_eagle_with_sliding_window():
1370
assert manager.block_pool.get_cached_block(
1371
block_hash_first_block, kv_cache_group_ids=[0]) is not None
1372
manager.block_pool.cached_block_hash_to_block.pop(
1354
- BlockHashWithGroupId(block_hash_first_block, 0))
1373
+ make_block_hash_with_group_id(block_hash_first_block, 0))
1374
1375
# New request
1376
req_after_evict = make_request("partial_eagle_after_evict", token_ids,
1358
- block_size, hash)
1377
+ block_size, sha256)
1378
computed_blocks, num_tokens = manager.get_computed_blocks(req_after_evict)
1379
# Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is
1380
# not considered. But after dropping the last matched block due to eagle,
tests/v1/core/test_single_type_kv_cache_manager.py
+8
-8
@@ -6,8 +6,8 @@ import random
6
import torch
7
8
from vllm.v1.core.block_pool import BlockPool
9
-from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
10
- KVCacheBlock)
9
+from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
10
+ make_block_hash_with_group_id)
11
from vllm.v1.core.single_type_kv_cache_manager import (
12
ChunkedLocalAttentionManager, SlidingWindowManager)
13
from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
@@ -44,7 +44,7 @@ def test_chunked_local_attention_possible_cached_prefix():
44
45
def run_one_case(block_is_cached, tail_token, expect_length):
46
block_hash_list = [
47
- BlockHash(i, ()) for i in range(len(block_is_cached))
47
+ BlockHash(str(i).encode()) for i in range(len(block_is_cached))
48
]
49
50
block_pool.cached_block_hash_to_block.clear()
@@ -53,8 +53,8 @@ def test_chunked_local_attention_possible_cached_prefix():
53
for i, (block_hash,
54
is_cached) in enumerate(zip(block_hash_list, block_is_cached)):
55
if is_cached:
56
- block_pool.cached_block_hash_to_block[BlockHashWithGroupId(
57
- block_hash, 0)] = {
56
+ block_pool.cached_block_hash_to_block[
57
+ make_block_hash_with_group_id(block_hash, 0)] = {
58
i: block_pool.blocks[i + 10],
59
}
60
@@ -109,7 +109,7 @@ def test_sliding_window_possible_cached_prefix():
109
110
def run_one_case(block_is_cached, expect_length):
111
block_hash_list = [
112
- BlockHash(i, ()) for i in range(len(block_is_cached))
112
+ BlockHash(str(i).encode()) for i in range(len(block_is_cached))
113
]
114
115
block_pool.cached_block_hash_to_block.clear()
@@ -118,8 +118,8 @@ def test_sliding_window_possible_cached_prefix():
118
for i, (block_hash,
119
is_cached) in enumerate(zip(block_hash_list, block_is_cached)):
120
if is_cached:
121
- block_pool.cached_block_hash_to_block[BlockHashWithGroupId(
122
- block_hash, 0)] = {
121
+ block_pool.cached_block_hash_to_block[
122
+ make_block_hash_with_group_id(block_hash, 0)] = {
123
i: block_pool.blocks[i + 10],
124
}
125
tests/v1/core/utils.py
+3
-2
@@ -9,6 +9,7 @@ from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,
9
from vllm.multimodal.inputs import (MultiModalFeatureSpec,
10
MultiModalKwargsItem, PlaceholderRange)
11
from vllm.sampling_params import SamplingParams
12
+from vllm.utils import sha256
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
@@ -130,10 +131,10 @@ def create_requests(
131
) -> list[Request]:
132
global _none_hash_initialized
133
if not _none_hash_initialized:
133
- init_none_hash(hash)
134
+ init_none_hash(sha256)
135
_none_hash_initialized = True
136
136
- block_hasher = get_request_block_hasher(block_size, hash)
137
+ block_hasher = get_request_block_hasher(block_size, sha256)
138
sampling_params = SamplingParams(ignore_eos=False,
139
max_tokens=max_tokens,
140
stop_token_ids=stop_token_ids,
tests/v1/engine/test_engine_args.py
+7
-6
@@ -36,18 +36,19 @@ def test_prefix_caching_from_cli():
36
assert vllm_config.cache_config.enable_prefix_caching
37
38
# default hash algorithm is "builtin"
39
- assert vllm_config.cache_config.prefix_caching_hash_algo == "builtin"
39
+ assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
40
+
41
+ # set hash algorithm to sha256_cbor
42
+ args = parser.parse_args(["--prefix-caching-hash-algo", "sha256_cbor"])
43
+ vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
44
+ assert vllm_config.cache_config.prefix_caching_hash_algo == \
45
+ "sha256_cbor"
46
47
# set hash algorithm to sha256
48
args = parser.parse_args(["--prefix-caching-hash-algo", "sha256"])
49
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
50
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
51
46
- # set hash algorithm to builtin
47
- args = parser.parse_args(["--prefix-caching-hash-algo", "builtin"])
48
- vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
49
- assert vllm_config.cache_config.prefix_caching_hash_algo == "builtin"
50
-
52
# an invalid hash algorithm raises an error
53
parser.exit_on_error = False
54
with pytest.raises(ArgumentError):
tests/v1/kv_connector/unit/utils.py
+3
-2
@@ -13,6 +13,7 @@ from vllm.distributed.kv_transfer.kv_connector.factory import (
13
KVConnectorFactory)
14
from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa
15
SharedStorageConnector)
16
+from vllm.utils import sha256
17
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
18
from vllm.v1.core.kv_cache_utils import (get_request_block_hasher,
19
init_none_hash)
@@ -127,11 +128,11 @@ def create_request(request_id: int,
128
use_all_1s_for_prompt_tokens: bool = False,
129
num_remote_blocks: int = 3,
130
block_size: int = 16,
130
- hash_fn: Callable = hash) -> Request:
131
+ hash_fn: Callable = sha256) -> Request:
132
"""Make dummy request for testing."""
133
global _none_hash_initialized
134
if not _none_hash_initialized:
134
- init_none_hash(hash)
135
+ init_none_hash(hash_fn)
136
_none_hash_initialized = True
137
138
kv_transfer_params: Optional[dict[str, Any]] = None
vllm/config/cache.py
+6
-11
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
24
BlockSize = Literal[1, 8, 16, 32, 64, 128]
25
CacheDType = Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2", "fp8_inc"]
26
MambaDType = Literal["auto", "float32"]
27
-PrefixCachingHashAlgo = Literal["builtin", "sha256", "sha256_cbor_64bit"]
27
+PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor"]
28
29
30
@config
@@ -63,17 +63,12 @@ class CacheConfig:
63
"""Sliding window size for the KV cache. This is primarily set in
64
`ModelConfig` and that value should be manually duplicated here."""
65
enable_prefix_caching: Optional[bool] = None
66
- """Whether to enable prefix caching. Disabled by default for V0. Enabled by
67
- default for V1."""
68
- prefix_caching_hash_algo: PrefixCachingHashAlgo = "builtin"
66
+ """Whether to enable prefix caching. Enabled by default for V1."""
67
+ prefix_caching_hash_algo: PrefixCachingHashAlgo = "sha256"
68
"""Set the hash algorithm for prefix caching:\n
70
- - "builtin" is Python's built-in hash.\n
71
- - "sha256" is collision resistant but with certain overheads.
72
- This option uses Pickle for object serialization before hashing.\n
73
- - "sha256_cbor_64bit" provides a reproducible, cross-language compatible
74
- hash. It serializes objects using canonical CBOR and hashes them with
75
- SHA-256. The resulting hash consists of the lower 64 bits of the SHA-256
76
- digest."""
69
+ - "sha256" uses Pickle for object serialization before hashing.\n
70
+ - "sha256_cbor" provides a reproducible, cross-language compatible hash. It
71
+ serializes objects using canonical CBOR and hashes them with SHA-256."""
72
cpu_offload_gb: float = 0
73
"""The space in GiB to offload to CPU, per GPU. Default is 0, which means
74
no offloading. Intuitively, this argument can be seen as a virtual way to
vllm/distributed/kv_events.py
+4
-3
@@ -16,6 +16,7 @@ import zmq
16
17
from vllm.config.kv_events import KVEventsConfig
18
from vllm.logger import init_logger
19
+from vllm.v1.core.kv_cache_utils import ExternalBlockHash
20
21
logger = init_logger(__name__)
22
@@ -44,8 +45,8 @@ MEDIUM_GPU = "GPU"
45
46
47
class BlockStored(KVCacheEvent):
47
- block_hashes: list[int]
48
- parent_block_hash: Optional[int]
48
+ block_hashes: list[ExternalBlockHash]
49
+ parent_block_hash: Optional[ExternalBlockHash]
50
token_ids: list[int]
51
block_size: int
52
lora_id: Optional[int]
@@ -53,7 +54,7 @@ class BlockStored(KVCacheEvent):
54
55
56
class BlockRemoved(KVCacheEvent):
56
- block_hashes: list[int]
57
+ block_hashes: list[ExternalBlockHash]
58
medium: Optional[str]
59
60
vllm/engine/arg_utils.py
+6
-14
@@ -1592,20 +1592,12 @@ class EngineArgs:
1592
"in low performance due to small KV cache size. Consider "
1593
"setting --max-model-len to a smaller value.", max_model_len)
1594
1595
- # if using prefix caching, we must set a hash algo
1596
- if self.enable_prefix_caching:
1597
- # Disable prefix caching for multimodal models for VLLM_V0.
1598
- if model_config.is_multimodal_model:
1599
- logger.warning(
1600
- "--enable-prefix-caching is not supported for multimodal "
1601
- "models in V0 and has been disabled.")
1602
- self.enable_prefix_caching = False
1603
-
1604
- # VLLM_V0 only supports builtin hash algo for prefix caching.
1605
- if self.prefix_caching_hash_algo == "sha256":
1606
- raise ValueError(
1607
- "sha256 is not supported for prefix caching in V0 engine. "
1608
- "Please use 'builtin'.")
1595
+ # Disable prefix caching for multimodal models for VLLM_V0.
1596
+ if self.enable_prefix_caching and model_config.is_multimodal_model:
1597
+ logger.warning(
1598
+ "--enable-prefix-caching is not supported for multimodal "
1599
+ "models in V0 and has been disabled.")
1600
+ self.enable_prefix_caching = False
1601
1602
# Set max_num_seqs to 256 for VLLM_V0.
1603
if self.max_num_seqs is None:
vllm/envs.py
+6
@@ -171,6 +171,7 @@ if TYPE_CHECKING:
171
VLLM_GPT_OSS_USE_CONTAINER_TOOL: bool = False
172
VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False
173
VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False
174
+ VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True
175
176
177
def get_default_cache_root():
@@ -1215,6 +1216,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1216
# Add optional custom scopes for profiling, disable to avoid overheads
1217
"VLLM_CUSTOM_SCOPES_FOR_PROFILING":
1218
lambda: bool(int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0"))),
1219
+
1220
+ # Represent block hashes in KV cache events as 64-bit integers instead of
1221
+ # raw bytes. Defaults to True for backward compatibility.
1222
+ "VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES":
1223
+ lambda: bool(int(os.getenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "1"))),
1224
}
1225
1226
# --8<-- [end:env-vars-definition]
vllm/utils/__init__.py
+10
-17
@@ -3249,7 +3249,7 @@ def check_use_alibi(model_config: ModelConfig) -> bool:
3249
and getattr(cfg.attn_config, "alibi", False)))))
3250
3251
3252
-def sha256(input) -> int:
3252
+def sha256(input) -> bytes:
3253
"""Hash any picklable Python object using SHA-256.
3254
3255
The input is serialized using pickle before hashing, which allows
@@ -3260,16 +3260,15 @@ def sha256(input) -> int:
3260
input: Any picklable Python object.
3261
3262
Returns:
3263
- An integer representing the SHA-256 hash of the serialized input.
3263
+ Bytes representing the SHA-256 hash of the serialized input.
3264
"""
3265
input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
3266
- return int.from_bytes(hashlib.sha256(input_bytes).digest(),
3267
- byteorder="big")
3266
+ return hashlib.sha256(input_bytes).digest()
3267
3268
3270
-def sha256_cbor_64bit(input) -> int:
3269
+def sha256_cbor(input) -> bytes:
3270
"""
3272
- Hash objects using CBOR serialization and SHA-256, then truncate to 64bits.
3271
+ Hash objects using CBOR serialization and SHA-256.
3272
3273
This option is useful for non-Python-dependent serialization and hashing.
3274
@@ -3280,17 +3279,13 @@ def sha256_cbor_64bit(input) -> int:
3279
Custom classes must implement CBOR serialization methods.
3280
3281
Returns:
3283
- An integer in the range [0, 2^64-1] representing the lower 64 bits
3284
- of the SHA-256 hash of the CBOR serialized input.
3282
+ Bytes representing the SHA-256 hash of the CBOR serialized input.
3283
"""
3284
input_bytes = cbor2.dumps(input, canonical=True)
3287
- full_hash = int.from_bytes(hashlib.sha256(input_bytes).digest(),
3288
- byteorder="big")
3285
+ return hashlib.sha256(input_bytes).digest()
3286
3290
- return full_hash & ((1 << 64) - 1)
3287
3292
-
3293
-def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], int]:
3288
+def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
3289
"""Get a hash function by name, or raise an error if
3290
the function is not found.
3291
Args:
@@ -3300,10 +3295,8 @@ def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], int]:
3295
"""
3296
if hash_fn_name == "sha256":
3297
return sha256
3303
- if hash_fn_name == "sha256_cbor_64bit":
3304
- return sha256_cbor_64bit
3305
- if hash_fn_name == "builtin":
3306
- return hash
3298
+ if hash_fn_name == "sha256_cbor":
3299
+ return sha256_cbor
3300
3301
raise ValueError(f"Unsupported hash function: {hash_fn_name}")
3302
vllm/v1/core/block_pool.py
+18
-9
@@ -9,7 +9,11 @@ from vllm.distributed.kv_events import (MEDIUM_GPU, AllBlocksCleared,
9
KVCacheEvent)
10
from vllm.logger import init_logger
11
from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
12
- FreeKVCacheBlockQueue, KVCacheBlock)
12
+ ExternalBlockHash,
13
+ FreeKVCacheBlockQueue, KVCacheBlock,
14
+ get_block_hash,
15
+ make_block_hash_with_group_id,
16
+ maybe_convert_block_hash)
17
from vllm.v1.request import Request
18
19
logger = init_logger(__name__)
@@ -84,8 +88,10 @@ class BlockPool:
88
"""
89
cached_blocks = []
90
for group_id in kv_cache_group_ids:
91
+ block_hash_with_group_id = make_block_hash_with_group_id(
92
+ block_hash, group_id)
93
cached_blocks_one_group = self.cached_block_hash_to_block.get(
88
- BlockHashWithGroupId(block_hash, group_id))
94
+ block_hash_with_group_id)
95
if not cached_blocks_one_group:
96
return None
97
first_block = next(iter(cached_blocks_one_group.values()))
@@ -124,28 +130,29 @@ class BlockPool:
130
assert len(request.block_hashes) >= num_full_blocks
131
new_block_hashes = request.block_hashes[num_cached_blocks:]
132
127
- new_hashes: Optional[list[int]] = ([] if self.enable_kv_cache_events
128
- else None)
133
+ new_hashes: Optional[list[ExternalBlockHash]] = (
134
+ [] if self.enable_kv_cache_events else None)
135
for i, blk in enumerate(new_full_blocks):
136
assert blk.block_hash is None
137
block_hash = new_block_hashes[i]
138
139
# Update and added the full block to the cache.
134
- block_hash_with_group_id = BlockHashWithGroupId(
140
+ block_hash_with_group_id = make_block_hash_with_group_id(
141
block_hash, kv_cache_group_id)
142
blk.block_hash = block_hash_with_group_id
143
self.cached_block_hash_to_block[block_hash_with_group_id][
144
blk.block_id] = blk
145
if new_hashes is not None:
140
- new_hashes.append(block_hash.hash_value)
146
+ new_hashes.append(maybe_convert_block_hash(block_hash))
147
148
if self.enable_kv_cache_events:
149
if num_cached_blocks == 0:
144
- parent_block_hash = None
150
+ parent_block_hash: Optional[ExternalBlockHash] = None
151
else:
152
parent_block = blocks[num_cached_blocks - 1]
153
assert parent_block.block_hash is not None
148
- parent_block_hash = parent_block.block_hash.get_hash_value()
154
+ parent_block_hash = maybe_convert_block_hash(
155
+ get_block_hash(parent_block.block_hash))
156
157
self.kv_event_queue.append(
158
BlockStored(
@@ -220,7 +227,9 @@ class BlockPool:
227
# we disable hybrid kv cache manager when kv cache event is
228
# enabled, so there is only one group.
229
self.kv_event_queue.append(
223
- BlockRemoved(block_hashes=[block_hash.get_hash_value()],
230
+ BlockRemoved(block_hashes=[
231
+ maybe_convert_block_hash(get_block_hash(block_hash))
232
+ ],
233
medium=MEDIUM_GPU))
234
return True
235
vllm/v1/core/kv_cache_utils.py
+64
-56
@@ -6,11 +6,12 @@ import os
6
from collections import defaultdict, deque
7
from collections.abc import Iterable, Sequence
8
from dataclasses import astuple, dataclass
9
-from typing import Any, Callable, NamedTuple, Optional
9
+from typing import Any, Callable, NewType, Optional, Union
10
11
+from vllm import envs
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
13
-from vllm.utils import GiB_bytes, cdiv, sha256_cbor_64bit
14
+from vllm.utils import GiB_bytes, cdiv, sha256_cbor
15
from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
16
FullAttentionSpec, KVCacheConfig,
17
KVCacheGroupSpec, KVCacheSpec,
@@ -18,59 +19,78 @@ from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
19
from vllm.v1.metrics.stats import PrefixCacheStats
20
from vllm.v1.request import Request
21
21
-logger = init_logger(__name__)
22
+# BlockHash represents the hash of a single KV-cache block used for
23
+# prefix caching. Treating it as a distinct type from ``bytes`` helps
24
+# catch accidental misuse when passing around raw byte strings.
25
+BlockHash = NewType("BlockHash", bytes)
26
+
27
+# ``BlockHashWithGroupId`` combines a ``BlockHash`` with its KV cache group ID.
28
+# It is represented as raw bytes for compactness and efficiency. The helper
29
+# functions below pack/unpack the ``BlockHash`` and group id into/from the key.
30
+BlockHashWithGroupId = NewType("BlockHashWithGroupId", bytes)
31
+
32
+# ExternalBlockHash is used for reproducible prefix-cache block hashing.
33
+# It's a union of ``bytes`` and ``int`` to keep backward compatibility
34
+# after we default block hashing to use sha256 bytes.
35
+ExternalBlockHash = Union[bytes, int]
36
37
24
-class BlockHash(NamedTuple):
25
- """Hash value of a block (int), the token IDs in the block, and extra keys.
26
- We keep a tuple of token IDs and extra keys to reduce the likelihood of
27
- hash collisions when the hash value is the same. By using SHA256 however,
28
- hash collisions are practically impossible.
38
+def make_block_hash_with_group_id(block_hash: BlockHash,
39
+ group_id: int) -> BlockHashWithGroupId:
40
+ """Pack a ``BlockHash`` and group id into a ``BlockHashWithGroupId``.
41
+
42
+ The group id is encoded using 4 bytes in big-endian order and appended to
43
+ the block hash bytes. This representation avoids creating tuples while
44
+ still allowing us to recover both components when needed.
45
"""
30
- # Hash value of the block in an integer.
31
- hash_value: int
32
- # Token IDs in the block.
33
- token_ids: tuple[int, ...]
34
- # Extra keys for the block.
35
- extra_keys: Optional[Any] = None
46
+ return BlockHashWithGroupId(block_hash +
47
+ group_id.to_bytes(4, "big", signed=False))
48
+
49
+
50
+def get_block_hash(key: BlockHashWithGroupId) -> BlockHash:
51
+ """Extract the ``BlockHash`` from a ``BlockHashWithGroupId``."""
52
+ return BlockHash(key[:-4])
53
54
38
-class BlockHashWithGroupId(NamedTuple):
39
- # The hash value for the contents (e.g., token_ids) of a block without group
40
- # ID. The value is the same for blocks representing the same tokens but for
41
- # different groups.
42
- block_hash: BlockHash
43
- # The KV cache group ID.
44
- group_id: int
55
+def get_group_id(key: BlockHashWithGroupId) -> int:
56
+ """Extract the group id from a ``BlockHashWithGroupId``."""
57
+ return int.from_bytes(key[-4:], "big", signed=False)
58
46
- def get_hash_value(self) -> int:
47
- return self.block_hash.hash_value
59
60
+def maybe_convert_block_hash(hash_bytes: BlockHash) -> ExternalBlockHash:
61
+ if not envs.VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES:
62
+ return hash_bytes
63
+ return int.from_bytes(hash_bytes, byteorder="big") & ((1 << 64) - 1)
64
+
65
+
66
+logger = init_logger(__name__)
67
68
# The hash seed for the first block of any prefix block sequence.
69
#
70
# We use a random value to avoid hash collisions or PYTHONHASHSEED environment
53
-# variable if set such that processes can share the seed if needed.
54
-# This aligns with the behavior of Python's hash() function, which also uses
55
-# a random seed if PYTHONHASHSEED is not set.
71
+# variable if set such that processes can share the seed if needed. This aligns
72
+# with the behavior of Python's hash() function, which also uses a random seed
73
+# if PYTHONHASHSEED is not set.
74
#
75
# The function `init_none_hash` initializes this variable globally.
58
-NONE_HASH: int
76
+NONE_HASH: BlockHash
77
78
61
-def init_none_hash(hash_fn: Callable):
79
+def init_none_hash(hash_fn: Callable[[Any], bytes]):
80
global NONE_HASH
81
82
hash_seed = os.getenv("PYTHONHASHSEED")
65
- if hash_seed is None and hash_fn is sha256_cbor_64bit:
83
+ if hash_seed is None and hash_fn is sha256_cbor:
84
logger.warning(
85
"PYTHONHASHSEED is not set. This will lead to non-reproducible "
68
- "block-hashes when using sha256_cbor_64bit as the hash function."
86
+ "block-hashes when using sha256_cbor as the hash function."
87
"Consider setting PYTHONHASHSEED to a fixed value for "
88
"reproducibility.")
89
72
- NONE_HASH = (int.from_bytes(os.urandom(32), byteorder="big")
73
- if hash_seed is None else hash_fn(hash_seed))
90
+ if hash_seed is None:
91
+ NONE_HASH = BlockHash(os.urandom(32))
92
+ else:
93
+ NONE_HASH = BlockHash(hash_fn(hash_seed))
94
95
96
class PrefixCachingMetrics:
@@ -142,8 +162,8 @@ class KVCacheBlock:
162
block_id: int
163
# Reference count.
164
ref_cnt: int = 0
145
- # The hash of the block composed of (block hash, tuple of token IDs).
146
- # It is only available when the block is full.
165
+ # The hash key (block hash + group id) of the block, only available
166
+ # when the block is full and cached.
167
_block_hash: Optional[BlockHashWithGroupId] = None
168
169
# Used to construct a doubly linked list for free blocks.
@@ -177,7 +197,7 @@ class KVCacheBlock:
197
if self.next_free_block else None)
198
return (f"KVCacheBlock(block_id={self.block_id}, "
199
f"ref_cnt={self.ref_cnt}, "
180
- f"_block_hash={self._block_hash}, "
200
+ f"_block_hash={self._block_hash!r}, "
201
f"prev_free_block={prev_block_id}, "
202
f"next_free_block={next_block_id})")
203
@@ -517,15 +537,14 @@ def generate_block_hash_extra_keys(
537
538
539
def hash_block_tokens(
520
- hash_function: Callable,
521
- parent_block_hash: Optional[int],
540
+ hash_function: Callable[[Any], bytes],
541
+ parent_block_hash: Optional[BlockHash],
542
curr_block_token_ids: Sequence[int],
543
extra_keys: Optional[tuple[Any, ...]] = None) -> BlockHash:
544
"""Computes a hash value corresponding to the contents of a block and
545
the contents of the preceding block(s). The hash value is used for
546
prefix caching. We use LRU cache for this function to avoid recomputing
547
hash values for the same block contents.
528
-
548
Args:
549
hash_function: The hash function used to compute block hash.
550
parent_block_hash: The hash of the parent block. None
@@ -533,7 +552,6 @@ def hash_block_tokens(
552
curr_block_token_ids: A list of token ids in the current
553
block. The current block is assumed to be full.
554
extra_keys: Extra keys for the block.
536
-
555
Returns:
556
The hash value of the block and the token ids in the block.
557
The entire tuple is used as the hash key of the block.
@@ -544,26 +562,16 @@ def hash_block_tokens(
562
curr_block_token_ids_tuple = tuple(curr_block_token_ids)
563
return BlockHash(
564
hash_function(
547
- (parent_block_hash, curr_block_token_ids_tuple, extra_keys)),
548
- curr_block_token_ids_tuple, extra_keys)
565
+ (parent_block_hash, curr_block_token_ids_tuple, extra_keys)))
566
567
568
def get_request_block_hasher(
569
block_size: int,
553
- caching_hash_fn: Callable[[Any],
554
- int]) -> Callable[[Request], list[BlockHash]]:
570
+ caching_hash_fn: Callable[[Any], bytes],
571
+) -> Callable[[Request], list[BlockHash]]:
572
"""
573
Returns a function which computes the list of un-computed block hashes
557
- of a request.
558
-
559
- Each request holds a list of its block hashes (request.block_hashes).
560
- When a request is created, it calls the below function to compute
561
- the hashes of all full blocks of the request's initial tokens.
562
- The hashes are then stored in request.block_hashes.
563
- Later, whenever new tokens are appended to the request, it calls
564
- the below function again to compute any new full blocks of tokens.
565
- The returned new hashes are appended to request.block_hashes.
566
- """
574
+ of a request."""
575
576
def request_block_hasher(request: Request) -> list[BlockHash]:
577
start_token_idx = len(request.block_hashes) * block_size
@@ -577,8 +585,8 @@ def get_request_block_hasher(
585
# last mm input.
586
curr_mm_idx = -1
587
580
- prev_block_hash_value = request.block_hashes[-1].hash_value \
581
- if request.block_hashes else None
588
+ prev_block_hash_value = (request.block_hashes[-1]
589
+ if request.block_hashes else None)
590
new_block_hashes: list[BlockHash] = []
591
while True:
592
end_token_idx = start_token_idx + block_size
@@ -598,7 +606,7 @@ def get_request_block_hasher(
606
607
new_block_hashes.append(block_hash)
608
start_token_idx += block_size
601
- prev_block_hash_value = block_hash.hash_value
609
+ prev_block_hash_value = block_hash
610
611
return new_block_hashes
612