@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 27df5199d9

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 15, 17:15 UTC

Support SHA256 as hash function in prefix caching (#15297)

Signed-off-by: Marko Rosenmueller <5467316+dr75@users.noreply.github.com>

marko committed Mar 26, 2025 at 19:11 UTC 27df5199d99627e1eb101071c2155f888181bd64

12 files changed +214 -71

docs/source/design/v1/prefix_caching.md

+4 -3

@@ -15,12 +15,13 @@ Block 3: |<------------------ prefix -------------------->| |<--- block tokens -

15 In the example above, the KV cache in the first block can be uniquely identified with the token “A gentle breeze stirred”. The third block can be uniquely identified with the tokens in the block “laughed in the distance”, along with the prefix tokens “A gentle breeze stirred the leaves as children”. Therefore, we can build the block hash of `hash(tuple[components])`, where components are:

17 * Parent hash value: The hash value of the parent hash block.

18 -* Block tokens: A tuple of tokens in this block. The reason to include the exact tokens is to reduce potential hash value collision.

18 +* Block tokens: A tuple of tokens in this block. The reason to include the exact tokens is to reduce potential hash value collision.

19 * Extra hashes: Other values required to make this block unique, such as LoRA IDs and multi-modality input hashes (see the example below).

21 -Note 1: We only cache full blocks.

21 +> **Note 1:** We only cache full blocks.

23 -Note 2: The above hash key structure is not 100% collision free. Theoretically it’s still possible for the different prefix tokens to have the same hash value, but this should be nearly impossible to happen. Of course, contributions are welcome if you have an awesome idea to eliminate collusion entirely.

23 +> **Note 2:** The above hash key structure is not 100% collision free. Theoretically it’s still possible for the different prefix tokens to have the same hash value. To avoid any hash collisions **in a multi-tenant setup, we advise to use SHA256** as hash function instead of the default builtin hash.

24 +SHA256 is supported since vLLM v0.8.3 and must be enabled with a command line argument. It comes with a performance impact of about 100-200ns per token (~6ms for 50k tokens of context).

26 **A hashing example with multi-modality inputs**

27 In this example, we illustrate how prefix caching works with multi-modality inputs (e.g., images). Assuming we have a request with the following messages:

tests/test_utils.py

+22 -1

@@ -2,6 +2,8 @@

2 # ruff: noqa

4 import asyncio

5 +import hashlib

6 +import pickle

7 import socket

8 from collections.abc import AsyncIterator

9 from unittest.mock import patch

@@ -14,7 +16,8 @@ from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

16 from vllm.utils import (FlexibleArgumentParser, MemorySnapshot,

17 PlaceholderModule, StoreBoolean, bind_kv_cache,

18 deprecate_kwargs, get_open_port, memory_profiling,

17 - merge_async_iterators, supports_kw, swap_dict_values)

19 + merge_async_iterators, sha256, supports_kw,

20 + swap_dict_values)

22 from .utils import create_new_process_for_each_test, error_on_warning

@@ -476,3 +479,21 @@ def test_swap_dict_values(obj, key1, key2):

479 assert obj[key1] == original_obj[key2]

480 else:

481 assert key1 not in obj

482 +

483 +@pytest.mark.parametrize("input", [(), ("abc", ), (None, ),

484 + (None, bool, [1, 2, 3])])

485 +@pytest.mark.parametrize("output", [0, 1, 2])

486 +def test_sha256(input: tuple, output: int):

487 + hash = sha256(input)

488 + assert hash is not None

489 + assert isinstance(hash, int)

490 + assert hash != 0

491 +

492 + bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)

493 + assert hash == int.from_bytes(hashlib.sha256(bytes).digest(), byteorder="big")

494 +

495 + # hashing again, returns the same value

496 + assert hash == sha256(input)

497 +

498 + # hashing different input, returns different value

499 + assert hash != sha256(input + (1, ))

tests/v1/core/test_kv_cache_utils.py

+29 -13

@@ -5,8 +5,12 @@ import torch

6 from vllm.multimodal.inputs import MultiModalKwargs

7 from vllm.sampling_params import SamplingParams

8 -from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,

9 - KVCacheBlock, PrefixCachingMetrics,

8 +from vllm.utils import sha256

9 +# disable yapf here as it formats differently than isort such that both fail

10 +# yapf: disable

11 +from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType,

12 + FreeKVCacheBlockQueue, KVCacheBlock,

13 + PrefixCachingMetrics,

14 generate_block_hash_extra_keys,

15 hash_block_tokens,

16 hash_request_tokens,

@@ -16,6 +20,8 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

20 from vllm.v1.metrics.stats import PrefixCacheStats

21 from vllm.v1.request import Request

23 +# yapf: enable

24 +

26 def make_request(request_id,

27 prompt_token_ids,

@@ -40,6 +46,12 @@ def make_request(request_id,

46 )

49 +def test_none_hash():

50 + assert NONE_HASH is not None

51 + assert isinstance(NONE_HASH, int)

52 + assert NONE_HASH != 0

53 +

54 +

55 def test_kv_cache_block():

56 # Test KVCacheBlock initialization

57 block = KVCacheBlock(block_id=0)

@@ -190,21 +202,23 @@ def test_generate_block_hash_extra_keys_no_mm_inputs():

202 assert next_mm_idx == 0

203

204

193 -def test_hash_block_tokens():

205 +@pytest.mark.parametrize("hash_fn", [sha256, hash])

206 +def test_hash_block_tokens(hash_fn):

207 parent_block_hash = 123

208 curr_block_token_ids = (1, 2, 3)

209 extra_keys = ("key1", "key2")

210

198 - block_hash = hash_block_tokens(parent_block_hash, curr_block_token_ids,

199 - extra_keys)

211 + block_hash = hash_block_tokens(hash_fn, parent_block_hash,

212 + curr_block_token_ids, extra_keys)

213 assert isinstance(block_hash, BlockHashType)

201 - assert block_hash.hash_value == hash(

214 + assert block_hash.hash_value == hash_fn(

215 (parent_block_hash, curr_block_token_ids, extra_keys))

216 assert block_hash.token_ids == curr_block_token_ids

217 assert block_hash.extra_keys == extra_keys

218

219

207 -def test_hash_request_tokens():

220 +@pytest.mark.parametrize("hash_fn", [sha256, hash])

221 +def test_hash_request_tokens(hash_fn):

222 request = make_request(

223 request_id=0,

224 prompt_token_ids=[_ for _ in range(6)],

@@ -219,7 +233,7 @@ def test_hash_request_tokens():

233 )

234

235 block_size = 3

222 - block_hashes = hash_request_tokens(block_size, request)

236 + block_hashes = hash_request_tokens(hash_fn, block_size, request)

237

238 assert len(block_hashes) == 2

239 assert isinstance(block_hashes[0], BlockHashType)

@@ -234,7 +248,8 @@ def test_hash_request_tokens():

248 assert block_hashes[1].extra_keys == ("hash2", )

249

250

237 -def test_hash_tokens_different_mm_input():

251 +@pytest.mark.parametrize("hash_fn", [sha256, hash])

252 +def test_hash_tokens_different_mm_input(hash_fn):

253 request1 = make_request(

254 request_id=0,

255 prompt_token_ids=[_ for _ in range(6)],

@@ -260,13 +275,14 @@ def test_hash_tokens_different_mm_input():

275 mm_hashes=["hash3", "hash2"],

276 )

277 block_size = 3

263 - block_hashes1 = hash_request_tokens(block_size, request1)

264 - block_hashes2 = hash_request_tokens(block_size, request2)

278 + block_hashes1 = hash_request_tokens(hash_fn, block_size, request1)

279 + block_hashes2 = hash_request_tokens(hash_fn, block_size, request2)

280 assert block_hashes1[0] != block_hashes2[0]

281 assert block_hashes1[1] != block_hashes2[1]

282

283

269 -def test_hash_request_tokens_no_mm_inputs():

284 +@pytest.mark.parametrize("hash_fn", [sha256, hash])

285 +def test_hash_request_tokens_no_mm_inputs(hash_fn):

286 request = make_request(

287 request_id=0,

288 prompt_token_ids=[_ for _ in range(6)],

@@ -275,7 +291,7 @@ def test_hash_request_tokens_no_mm_inputs():

291 )

292

293 block_size = 3

278 - block_hashes = hash_request_tokens(block_size, request)

294 + block_hashes = hash_request_tokens(hash_fn, block_size, request)

295

296 assert len(block_hashes) == 2

297 assert block_hashes[0].token_ids == (0, 1, 2)

tests/v1/core/test_prefix_caching.py

+17 -5

@@ -7,7 +7,7 @@ import pytest

8 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

9 from vllm.sampling_params import SamplingParams

10 -from vllm.utils import cdiv

10 +from vllm.utils import cdiv, sha256

11 from vllm.v1.core.block_pool import BlockPool

12 from vllm.v1.core.kv_cache_manager import KVCacheManager, Request

13 from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,

@@ -39,16 +39,21 @@ def make_request(request_id,

39 )

42 -def test_prefill():

42 +@pytest.mark.parametrize("hash_algo", ["sha256", "hash"])

43 +def test_prefill(hash_algo):

44 manager = KVCacheManager(

45 block_size=16,

46 num_gpu_blocks=10,

47 max_model_len=8192,

48 sliding_window=None,

49 enable_caching=True,

50 + caching_hash_algo=hash_algo,

51 num_preallocate_tokens=16,

52 )

54 + # choose the hash function according to the parameter

55 + hash_fn = sha256 if hash_algo == "sha256" else hash

56 +

57 # Complete 3 blocks (48 tokens)

58 common_token_ids = [i for i in range(3) for _ in range(16)]

@@ -68,7 +73,8 @@ def test_prefill():

73 parent_block_hash = None

74 for block_id in (0, 1, 2):

75 block_tokens = tuple(all_token_ids[block_id * 16:(block_id + 1) * 16])

71 - block_hash = hash_block_tokens(parent_block_hash, block_tokens)

76 + block_hash = hash_block_tokens(hash_fn, parent_block_hash,

77 + block_tokens)

78 assert manager.block_pool.blocks[block_id].block_hash == block_hash

79 assert manager.block_pool.blocks[block_id].ref_cnt == 1

80 parent_block_hash = block_hash.hash_value

@@ -163,6 +169,8 @@ def test_prefill_plp():

169 enable_caching=True,

170 num_preallocate_tokens=16,

171 )

172 + # the default hash function is hash

173 + hash_fn = hash

174

175 # Complete 3 blocks (48 tokens)

176 common_token_ids = [i for i in range(3) for _ in range(16)]

@@ -185,7 +193,8 @@ def test_prefill_plp():

193 parent_block_hash = None

194 for block_id in (0, 1, 2):

195 block_tokens = tuple(all_token_ids[block_id * 16:(block_id + 1) * 16])

188 - block_hash = hash_block_tokens(parent_block_hash, block_tokens)

196 + block_hash = hash_block_tokens(hash_fn, parent_block_hash,

197 + block_tokens)

198 assert manager.block_pool.blocks[block_id].block_hash == block_hash

199 assert manager.block_pool.blocks[block_id].ref_cnt == 1

200 parent_block_hash = block_hash.hash_value

@@ -522,7 +531,8 @@ def test_preallocate_blocks(num_preallocate_tokens: int, block_size: int):

531 assert len(blocks) == 1 + num_preallocated_blocks

532

533

525 -def test_cache_blocks():

534 +@pytest.mark.parametrize("hash_fn", [sha256, hash])

535 +def test_cache_blocks(hash_fn):

536 """

537 This is a unit test that tests the correctness of the _cache_full_blocks

538 function of KVCacheManager.

@@ -550,6 +560,7 @@ def test_cache_blocks():

560 num_cached_blocks=0,

561 num_full_blocks=2,

562 block_size=block_size,

563 + hash_fn=hash_fn,

564 )

565

566 assert len(block_pool.cached_block_hash_to_block) == 2

@@ -564,6 +575,7 @@ def test_cache_blocks():

575 num_cached_blocks=2,

576 num_full_blocks=3,

577 block_size=block_size,

578 + hash_fn=hash_fn,

579 )

580 assert len(block_pool.cached_block_hash_to_block) == 3

581 assert blocks[0].block_hash is not None

tests/v1/engine/test_engine_args.py

+20

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 +from argparse import ArgumentError

4 +

5 import pytest

7 from vllm import envs

@@ -32,6 +34,24 @@ def test_prefix_caching_from_cli():

34 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()

35 assert vllm_config.cache_config.enable_prefix_caching

37 + # default hash algorithm is "builtin"

38 + assert vllm_config.cache_config.prefix_caching_hash_algo == "builtin"

39 +

40 + # set hash algorithm to sha256

41 + args = parser.parse_args(["--prefix-caching-hash-algo", "sha256"])

42 + vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()

43 + assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"

44 +

45 + # set hash algorithm to builtin

46 + args = parser.parse_args(["--prefix-caching-hash-algo", "builtin"])

47 + vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()

48 + assert vllm_config.cache_config.prefix_caching_hash_algo == "builtin"

49 +

50 + # an invalid hash algorithm raises an error

51 + parser.exit_on_error = False

52 + with pytest.raises(ArgumentError):

53 + args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])

54 +

56 def test_defaults_with_usage_context():

57 engine_args = EngineArgs(model="facebook/opt-125m")

vllm/config.py

@@ -1124,6 +1124,7 @@ class CacheConfig:

1124 num_gpu_blocks_override: Optional[int] = None,

1125 sliding_window: Optional[int] = None,

1126 enable_prefix_caching: bool = False,

1127 + prefix_caching_hash_algo: str = "builtin",

1128 cpu_offload_gb: float = 0,

1129 calculate_kv_scales: Optional[bool] = None,

1130 ) -> None:

@@ -1135,6 +1136,7 @@ class CacheConfig:

1136 self.is_attention_free = is_attention_free

1137 self.sliding_window = sliding_window

1138 self.enable_prefix_caching = enable_prefix_caching

1139 + self.prefix_caching_hash_algo = prefix_caching_hash_algo

1140 self.cpu_offload_gb = cpu_offload_gb

1141 self.calculate_kv_scales = calculate_kv_scales

1142 self._verify_args()

@@ -1185,6 +1187,13 @@ class CacheConfig:

1187 "Prefix caching is not supported with sliding window. "

1188 "Run with --disable-sliding-window to use prefix caching.")

1189

1190 + if self.enable_prefix_caching and self.prefix_caching_hash_algo not in (

1191 + "builtin", "sha256"):

1192 + raise ValueError(

1193 + "Unknown prefix caching hash algorithm: "

1194 + f"{self.prefix_caching_hash_algo}. Must be either "

1195 + "'builtin' or 'sha256'.")

1196 +

1197 def verify_with_parallel_config(

1198 self,

1199 parallel_config: "ParallelConfig",

vllm/engine/arg_utils.py

+32 -6

@@ -118,6 +118,7 @@ class EngineArgs:

118 max_parallel_loading_workers: Optional[int] = None

119 block_size: Optional[int] = None

120 enable_prefix_caching: Optional[bool] = None

121 + prefix_caching_hash_algo: str = "builtin"

122 disable_sliding_window: bool = False

123 disable_cascade_attn: bool = False

124 use_v2_block_manager: bool = True

@@ -475,6 +476,16 @@ class EngineArgs:

476 help="Enables automatic prefix caching. "

477 "Use ``--no-enable-prefix-caching`` to disable explicitly.",

478 )

479 + parser.add_argument(

480 + "--prefix-caching-hash-algo",

481 + type=str,

482 + choices=["builtin", "sha256"],

483 + default=EngineArgs.prefix_caching_hash_algo,

484 + help="Set the hash algorithm for prefix caching. "

485 + "Options are 'builtin' (Python's built-in hash) or 'sha256' "

486 + "(collision resistant but with certain overheads). Defaults "

487 + "to 'builtin'.",

488 + )

489 parser.add_argument('--disable-sliding-window',

490 action='store_true',

491 help='Disables sliding window, '

@@ -1329,6 +1340,7 @@ class EngineArgs:

1340 num_gpu_blocks_override=self.num_gpu_blocks_override,

1341 sliding_window=model_config.get_sliding_window(),

1342 enable_prefix_caching=self.enable_prefix_caching,

1343 + prefix_caching_hash_algo=self.prefix_caching_hash_algo,

1344 cpu_offload_gb=self.cpu_offload_gb,

1345 calculate_kv_scales=self.calculate_kv_scales,

1346 )

@@ -1737,12 +1749,22 @@ class EngineArgs:

1749 msg = "Chunked prefill is not supported for pooling models"

1750 raise ValueError(msg)

1751

1740 - # Disable prefix caching for multimodal models for VLLM_V0.

1741 - if (model_config.is_multimodal_model and self.enable_prefix_caching):

1742 - logger.warning(

1743 - "--enable-prefix-caching is not supported for multimodal "

1744 - "models in V0 and has been disabled.")

1745 - self.enable_prefix_caching = False

1752 + # if using prefix caching, we must set a hash algo

1753 + if self.enable_prefix_caching:

1754 + # Disable prefix caching for multimodal models for VLLM_V0.

1755 + if model_config.is_multimodal_model:

1756 + logger.warning(

1757 + "--enable-prefix-caching is not supported for multimodal "

1758 + "models in V0 and has been disabled.")

1759 + self.enable_prefix_caching = False

1760 +

1761 + # VLLM_V0 only supports builtin hash algo for prefix caching.

1762 + if self.prefix_caching_hash_algo is None:

1763 + self.prefix_caching_hash_algo = "builtin"

1764 + elif self.prefix_caching_hash_algo == "sha256":

1765 + raise ValueError(

1766 + "sha256 is not supported for prefix caching in V0 engine. "

1767 + "Please use 'builtin'.")

1768

1769 # Set max_num_seqs to 256 for VLLM_V0.

1770 if self.max_num_seqs is None:

@@ -1758,6 +1780,10 @@ class EngineArgs:

1780 if self.enable_prefix_caching is None:

1781 self.enable_prefix_caching = True

1782

1783 + # if using prefix caching, we must set a hash algo

1784 + if self.enable_prefix_caching and self.prefix_caching_hash_algo is None:

1785 + self.prefix_caching_hash_algo = "builtin"

1786 +

1787 # V1 should use the new scheduler by default.

1788 # Swap it only if this arg is set to the original V0 default

1789 if self.scheduler_cls == EngineArgs.scheduler_cls:

vllm/utils.py

+20

@@ -10,6 +10,7 @@ import datetime

10 import enum

11 import gc

12 import getpass

13 +import hashlib

14 import importlib

15 import importlib.metadata

16 import importlib.util

@@ -17,6 +18,7 @@ import inspect

18 import ipaddress

19 import multiprocessing

20 import os

21 +import pickle

22 import re

23 import signal

24 import socket

@@ -2442,3 +2444,21 @@ def cprofile(save_file: Optional[str] = None, enabled: bool = True):

2444 return wrapper

2445

2446 return decorator

2447 +

2448 +

2449 +def sha256(input) -> int:

2450 + """Hash any picklable Python object using SHA-256.

2451 +

2452 + The input is serialized using pickle before hashing, which allows

2453 + arbitrary Python objects to be used. Note that this function does

2454 + not use a hash seed—if you need one, prepend it explicitly to the input.

2455 +

2456 + Args:

2457 + input: Any picklable Python object.

2458 +

2459 + Returns:

2460 + An integer representing the SHA-256 hash of the serialized input.

2461 + """

2462 + input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)

2463 + return int.from_bytes(hashlib.sha256(input_bytes).digest(),

2464 + byteorder="big")

vllm/v1/core/block_pool.py

+11 -9

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 from collections import defaultdict

3 from collections.abc import Iterable

4 -from typing import Optional

4 +from typing import Callable, Optional

6 from vllm.logger import init_logger

7 from vllm.v1.core.kv_cache_utils import (BlockHashType, FreeKVCacheBlockQueue,

@@ -15,10 +15,10 @@ logger = init_logger(__name__)

16 class BlockPool:

17 """BlockPool that manages KVCacheBlocks.

18 - It provides methods to allocate, free and cache the kv cache blocks. The

19 - free_block_queue stores the free blocks in eviction order to enable

20 - allocation, free, and cache eviction. The cached_block_hash_to_block

21 - maps between block hash and cached block to support finding cached blocks

18 + It provides methods to allocate, free and cache the kv cache blocks. The

19 + free_block_queue stores the free blocks in eviction order to enable

20 + allocation, free, and cache eviction. The cached_block_hash_to_block

21 + maps between block hash and cached block to support finding cached blocks

22 by their block hash.

24 Args:

@@ -75,11 +75,12 @@ class BlockPool:

75 num_cached_blocks: int,

76 num_full_blocks: int,

77 block_size: int,

78 + hash_fn: Callable,

79 ) -> None:

80 """Cache a list of full blocks for prefix caching.

81 This function takes a list of blocks that will have their block hash

82 metadata to be updated and cached. Given a request, it computes the

82 - block hashes for the blocks starting from `num_cached_blocks` to

83 + block hashes for the blocks starting from `num_cached_blocks` to

84 `num_full_blocks`, updating the metadata for each block

85 and caching them in the `cached_block_hash_to_block`.

@@ -87,12 +88,13 @@ class BlockPool:

88 request: The request to cache the blocks.

89 blocks: All blocks in the request.

90 block_hashes: Block hashes of the blocks in the request. Note that

90 - this list may be shorter than the blocks list. In this case the

91 + this list may be shorter than the blocks list. In this case the

92 missed block hash will be computed in this function.

93 num_cached_blocks: The number of blocks that are already cached.

93 - num_full_blocks: The number of blocks that are full and should

94 + num_full_blocks: The number of blocks that are full and should

95 be cached after this function.

96 block_size: Number of tokens in each block.

97 + hash_fn: The hash function to use for block hashes.

98 """

99 if num_cached_blocks == num_full_blocks:

100 return

@@ -138,7 +140,7 @@ class BlockPool:

140 request, start_token_idx, end_token_idx, -1)

141

142 # Compute the hash of the current block.

141 - block_hash = hash_block_tokens(prev_block_hash_value,

143 + block_hash = hash_block_tokens(hash_fn, prev_block_hash_value,

144 block_tokens, extra_keys)

145 block_hashes.append(block_hash)

146

vllm/v1/core/kv_cache_manager.py

+6 -2

@@ -5,7 +5,7 @@ from collections.abc import Iterable

5 from typing import Optional

7 from vllm.logger import init_logger

8 -from vllm.utils import cdiv

8 +from vllm.utils import cdiv, sha256

9 from vllm.v1.core.block_pool import BlockPool

10 from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,

11 hash_request_tokens)

@@ -24,6 +24,7 @@ class KVCacheManager:

24 max_model_len: int,

25 sliding_window: Optional[int] = None,

26 enable_caching: bool = True,

27 + caching_hash_algo: str = "builtin",

28 num_preallocate_tokens: int = 64,

29 log_stats: bool = False,

30 ) -> None:

@@ -33,6 +34,7 @@ class KVCacheManager:

34 self.max_num_blocks_per_req = cdiv(max_model_len, block_size)

35 self.sliding_window = sliding_window

36 self.enable_caching = enable_caching

37 + self.caching_hash_fn = sha256 if caching_hash_algo == "sha256" else hash

38 # FIXME: make prefix cache stats conditional on log_stats

39 self.log_stats = log_stats

40 # NOTE(woosuk): To avoid frequent block allocation, we preallocate some

@@ -109,7 +111,8 @@ class KVCacheManager:

111 # if the scheduler has tried to schedule the request before.

112 block_hashes = self.req_to_block_hashes[request.request_id]

113 if not block_hashes:

112 - block_hashes = hash_request_tokens(self.block_size, request)

114 + block_hashes = hash_request_tokens(self.caching_hash_fn,

115 + self.block_size, request)

116 self.req_to_block_hashes[request.request_id] = block_hashes

117

118 self.prefix_cache_stats.requests += 1

@@ -247,6 +250,7 @@ class KVCacheManager:

250 num_cached_blocks=num_cached_blocks,

251 num_full_blocks=num_full_blocks_after_append,

252 block_size=self.block_size,

253 + hash_fn=self.caching_hash_fn,

254 )

255

256 self.num_cached_block[

vllm/v1/core/kv_cache_utils.py

+43 -32

@@ -1,12 +1,14 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """KV-Cache Utilities."""

3 +import os

4 from collections import deque

5 from collections.abc import Sequence

6 from dataclasses import dataclass

6 -from typing import Any, NamedTuple, Optional

7 +from typing import Any, Callable, NamedTuple, Optional

9 from vllm.config import VllmConfig

10 from vllm.logger import init_logger

11 +from vllm.utils import sha256

12 from vllm.v1.kv_cache_interface import (KVCacheConfig, KVCacheGroupSpec,

13 KVCacheSpec, KVCacheTensor)

14 from vllm.v1.metrics.stats import PrefixCacheStats

@@ -18,9 +20,8 @@ logger = init_logger(__name__)

20 class BlockHashType(NamedTuple):

21 """Hash value of a block (int), the token IDs in the block, and extra keys.

22 We keep a tuple of token IDs and extra keys to reduce the likelihood of

21 - hash collisions when the hash value is the same. But please note that

22 - hash collisions can still theoretically occur, albeit with an extremely

23 - low probability.

23 + hash collisions when the hash value is the same. By using SHA256 however,

24 + hash collisions are practically impossible.

25 """

26 # Hash value of the block in an integer.

27 hash_value: int

@@ -30,6 +31,20 @@ class BlockHashType(NamedTuple):

31 extra_keys: Optional[Any] = None

34 +# The hash seed for the first block of the prefix block sequence.

35 +#

36 +# Even if the hash function is the builtin hash(), we use sha256 to generate

37 +# the initial hash to simplify the code. This is not performance critical

38 +# as it is done one per process.

39 +#

40 +# We use a random value to avoid hash collisions or PYTHONHASHSEED environment

41 +# variable if set such that processes can share the seed if needed.

42 +# This aligns with the behavior of Python's hash() function, which also uses

43 +# a random seed if PYTHONHASHSEED is not set.

44 +NONE_HASH = int.from_bytes(os.urandom(32), byteorder="big") if os.getenv(

45 + 'PYTHONHASHSEED') is not None else sha256(os.getenv('PYTHONHASHSEED'))

46 +

47 +

48 class PrefixCachingMetrics:

49 """Metrics for prefix caching with a hit rate of the most recent N requests.

@@ -148,7 +163,7 @@ class FreeKVCacheBlockQueue:

163 builtin deque to support removing a block in the middle of the queue

164 in O(1) time. To close the performance gap to the builtin deque which is

165 implemented in C++, this class does not allocate any Python objects when

151 - manipulating the linked list. Instead, this class manipulates the

166 + manipulating the linked list. Instead, this class manipulates the

167 prev_free_block and next_free_block attributes of the given blocks.

168

169 The queue is ordered by block ID in the beginning. When a block is allocated

@@ -178,7 +193,7 @@ class FreeKVCacheBlockQueue:

193

194 def popleft(self) -> KVCacheBlock:

195 """Pop the first free block and reduce num_free_blocks by 1.

181 -

196 +

197 Returns:

198 The first free block.

199 """

@@ -191,7 +206,7 @@ class FreeKVCacheBlockQueue:

206

207 def remove(self, block: KVCacheBlock) -> None:

208 """Remove a block in the free list and reduce num_free_blocks by 1.

194 -

209 +

210 Args:

211 block: The block to remove.

212 """

@@ -235,7 +250,7 @@ class FreeKVCacheBlockQueue:

250

251 def get_all_free_blocks(self) -> list[KVCacheBlock]:

252 """Get all free blocks in the free list. Mainly used for testing.

238 -

253 +

254 Returns:

255 A list of free blocks.

256 """

@@ -251,10 +266,10 @@ def need_extra_keys(request: Request) -> bool:

266 """Check whether the blocks allocated to this request need extra hash keys.

267

268 Args:

254 - request (Request): The request.

269 + request (Request): The request.

270

271 Returns:

257 - bool: Whether blocks allocated to this request need extra hash keys.

272 + bool: Whether blocks allocated to this request need extra hash keys.

273 """

274

275 # Multimodal requests need to include the MM hash.

@@ -269,13 +284,13 @@ def _gen_mm_extra_hash_keys(request: Request, start_token_idx: int,

284 computation. For multi-modal inputs, the extra keys are

285 (mm_hash, start_offset) that indicate a mm input contained in the

286 block and its starting offset in the block tokens.

272 -

287 +

288 Args:

289 request: The request object.

290 start_token_idx: The start token index of the block.

291 end_token_idx: The end token index of the block.

292 start_mm_idx: The start multi-modal index of the block.

278 -

293 +

294 Returns:

295 A tuple of extra keys and the next multi-modal index.

296 """

@@ -333,10 +348,10 @@ def _gen_mm_extra_hash_keys(request: Request, start_token_idx: int,

348

349 def _gen_lora_extra_hash_keys(request: Request) -> list[int]:

350 """Generate extra keys related to LoRA for block hash computation.

336 -

351 +

352 Args:

353 request: The request object.

339 -

354 +

355 Returns:

356 Return LoRA id of the request if it is a LoRA request. Return empty

357 list otherwise.

@@ -351,13 +366,13 @@ def generate_block_hash_extra_keys(

366 start_mm_idx: int) -> tuple[Optional[tuple[Any, ...]], int]:

367 """Generate extra keys for the block hash. The extra keys can come from

368 the multi-modal inputs and request specific metadata (e.g., LoRA ID).

354 -

369 +

370 Args:

371 request: The request object.

372 start_token_idx: The start token index of the block.

373 end_token_idx: The end token index of the block.

374 start_mm_idx: The start multi-modal index of the block.

360 -

375 +

376 Returns:

377 A tuple of extra keys and the next multi-modal index.

378 """

@@ -375,6 +390,7 @@ def generate_block_hash_extra_keys(

390

391

392 def hash_block_tokens(

393 + hash_function: Callable,

394 parent_block_hash: Optional[int],

395 curr_block_token_ids: Sequence[int],

396 extra_keys: Optional[tuple[Any, ...]] = None) -> BlockHashType:

@@ -395,21 +411,16 @@ def hash_block_tokens(

411 The entire tuple is used as the hash key of the block.

412 """

413 if not parent_block_hash:

398 - # Note that we use 'None' as a string here instead of None because

399 - # as of Python 3.12, hash(None) returns a constant predictable value.

400 - # This could possibly make it easier to find and exploit hash

401 - # collisions. 'None' as a string will be hashed differently per process,

402 - # but consistently within the same process. This is the same as the

403 - # behavior of None prior to Python 3.12.

404 - parent_block_hash = hash('None')

414 + parent_block_hash = NONE_HASH

415

416 curr_block_token_ids_tuple = tuple(curr_block_token_ids)

417 return BlockHashType(

408 - hash((parent_block_hash, curr_block_token_ids_tuple, extra_keys)),

418 + hash_function(

419 + (parent_block_hash, curr_block_token_ids_tuple, extra_keys)),

420 curr_block_token_ids_tuple, extra_keys)

421

422

412 -def hash_request_tokens(block_size: int,

423 +def hash_request_tokens(hash_function: Any, block_size: int,

424 request: Request) -> list[BlockHashType]:

425 """Computes hash values of a chain of blocks given a sequence of

426 token IDs. The hash value is used for prefix caching.

@@ -441,7 +452,7 @@ def hash_request_tokens(block_size: int,

452 req_extra_keys, curr_mm_idx = generate_block_hash_extra_keys(

453 request, start, end, curr_mm_idx)

454

444 - block_hash = hash_block_tokens(parent_block_hash_value,

455 + block_hash = hash_block_tokens(hash_function, parent_block_hash_value,

456 block_token_ids, req_extra_keys)

457 ret.append(block_hash)

458 parent_block_hash_value = block_hash.hash_value

@@ -452,7 +463,7 @@ def check_enough_kv_cache_memory(vllm_config: VllmConfig,

463 kv_cache_spec: dict[str, KVCacheSpec],

464 available_memory: int):

465 """

455 - Checks whether `available_memory` is enough for the KV cache to hold at

466 + Checks whether `available_memory` is enough for the KV cache to hold at

467 least one request with the model's max_model_len.

468

469 Args:

@@ -489,15 +500,15 @@ def create_kv_cache_group_specs(

500 grouped_layer_names: list[list[str]]) -> list[KVCacheGroupSpec]:

501 """

502 Create KVCacheGroupSpec object for each kv cache group layer.

492 - The layers in the same group should share the same

503 + The layers in the same group should share the same

504 KVCacheSpec.

505

506 Args:

507 kv_cache_spec:

508 A mapping from each layer name to its corresponding KVCacheSpec.

509 grouped_layer_names:

499 - A list of kv cache groups, where each element is a list of layer

500 - names that belong to the same group and should share the same

510 + A list of kv cache groups, where each element is a list of layer

511 + names that belong to the same group and should share the same

512 KVCacheSpec.

513 Returns:

514 A list of KVCacheGroupSpec objects, one for each group.

@@ -614,11 +625,11 @@ def get_kv_cache_config(vllm_config: VllmConfig,

625

626 def unify_kv_cache_configs(kv_cache_configs: list[KVCacheConfig]):

627 """

617 - Make the KV cache configurations for each worker consistent, so that all

628 + Make the KV cache configurations for each worker consistent, so that all

629 workers can be controlled by the same KVCacheManager.

630 This function verifies that the layer group of each worker are the same,

631 and changes the num_blocks of each worker to the smallest among all workers.

621 -

632 +

633 Args:

634 kv_cache_configs: The KV cache configurations for each worker. Will be

635 in-place modified to make them consistent.

vllm/v1/core/sched/scheduler.py

@@ -61,6 +61,7 @@ class Scheduler(SchedulerInterface):

61 max_model_len=self.max_model_len,

62 sliding_window=self.cache_config.sliding_window,

63 enable_caching=self.cache_config.enable_prefix_caching,

64 + caching_hash_algo=self.cache_config.prefix_caching_hash_algo,

65 log_stats=self.log_stats)

66 self.block_size = self.cache_config.block_size