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, 09:45 UTC
[UT] Add ut for none hash (#17892)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
May 15, 2025 at 13:28 UTC
420caf7557f85635fae09b3f2c27e38ac751551f
1 file changed
+31
-10
tests/v1/core/test_kv_cache_utils.py
+31
-10
@@ -1,4 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import importlib
3
4
import pytest
5
import torch
@@ -10,8 +11,7 @@ from vllm.utils import GiB_bytes, sha256
11
from vllm.v1.core.kv_cache_manager import KVCacheManager
12
# disable yapf here as it formats differently than isort such that both fail
13
# yapf: disable
13
-from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType,
14
- FreeKVCacheBlockQueue, KVCacheBlock,
14
+from vllm.v1.core.kv_cache_utils import (FreeKVCacheBlockQueue, KVCacheBlock,
15
PrefixCachingMetrics,
16
estimate_max_model_len,
17
generate_block_hash_extra_keys,
@@ -65,13 +65,29 @@ def new_kv_cache_spec(block_size=16,
65
sliding_window=sliding_window)
66
67
68
-def test_none_hash():
69
- assert NONE_HASH is not None
70
- assert isinstance(NONE_HASH, int)
71
- assert NONE_HASH != 0
68
+def test_none_hash(monkeypatch):
69
+ import vllm.v1.core.kv_cache_utils
70
+
71
+ # case 1: PYTHONHASHSEED is not set, use random
72
+ with monkeypatch.context() as m:
73
+ m.delenv('PYTHONHASHSEED', raising=False)
74
+ reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
75
+ assert reloaded_kv_cache_utils.NONE_HASH is not None
76
+ assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
77
+ assert reloaded_kv_cache_utils.NONE_HASH != 0
78
+
79
+ # case 2: PYTHONHASHSEED is set, use the seed
80
+ with monkeypatch.context() as m:
81
+ m.setenv('PYTHONHASHSEED', 'python hash seed')
82
+ reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
83
+ assert reloaded_kv_cache_utils.NONE_HASH is not None
84
+ assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
85
+ assert sha256('python hash seed') == reloaded_kv_cache_utils.NONE_HASH
86
87
88
def test_kv_cache_block():
89
+ import vllm.v1.core.kv_cache_utils
90
+
91
# Test KVCacheBlock initialization
92
block = KVCacheBlock(block_id=0)
93
assert block.block_id == 0
@@ -85,7 +101,8 @@ def test_kv_cache_block():
101
assert block.ref_cnt == 0
102
103
# Test block hash setting and resetting
88
- block_hash = BlockHashType(hash_value=123, token_ids=(1, 2, 3))
104
+ block_hash = vllm.v1.core.kv_cache_utils.BlockHashType(hash_value=123,
105
+ token_ids=(1, 2, 3))
106
block.block_hash = block_hash
107
assert block.block_hash == block_hash
108
@@ -259,13 +276,14 @@ def test_generate_block_hash_extra_keys_cache_salt():
276
277
@pytest.mark.parametrize("hash_fn", [sha256, hash])
278
def test_hash_block_tokens(hash_fn):
279
+ import vllm.v1.core.kv_cache_utils
280
parent_block_hash = 123
281
curr_block_token_ids = (1, 2, 3)
282
extra_keys = ("key1", "key2")
283
284
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
285
curr_block_token_ids, extra_keys)
268
- assert isinstance(block_hash, BlockHashType)
286
+ assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHashType)
287
assert block_hash.hash_value == hash_fn(
288
(parent_block_hash, curr_block_token_ids, extra_keys))
289
assert block_hash.token_ids == curr_block_token_ids
@@ -274,6 +292,7 @@ def test_hash_block_tokens(hash_fn):
292
293
@pytest.mark.parametrize("hash_fn", [sha256, hash])
294
def test_hash_request_tokens(hash_fn):
295
+ import vllm.v1.core.kv_cache_utils
296
request = make_request(
297
request_id=0,
298
prompt_token_ids=[_ for _ in range(6)],
@@ -288,8 +307,10 @@ def test_hash_request_tokens(hash_fn):
307
block_hashes = hash_request_tokens(hash_fn, block_size, request)
308
309
assert len(block_hashes) == 2
291
- assert isinstance(block_hashes[0], BlockHashType)
292
- assert isinstance(block_hashes[1], BlockHashType)
310
+ assert isinstance(block_hashes[0],
311
+ vllm.v1.core.kv_cache_utils.BlockHashType)
312
+ assert isinstance(block_hashes[1],
313
+ vllm.v1.core.kv_cache_utils.BlockHashType)
314
315
# Check the first block
316
assert block_hashes[0].token_ids == (0, 1, 2)