xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 08:34:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[UT] Add ut for none hash (#17892)

Browse Source 
Signed-off-by: Andy Xie <andy.xning@gmail.com>
This commit is contained in:
Ning Xie 2025-05-15 13:28:11 +08:00 committed by GitHub
parent 4f07a64075
commit 420caf7557
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 31 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
tests/v1/core/test_kv_cache_utils.py
View File

@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
import importlib
import pytest import pytest
import torch import torch
@ -10,8 +11,7 @@ from vllm.utils import GiB_bytes, sha256
from vllm.v1.core.kv_cache_manager import KVCacheManager from vllm.v1.core.kv_cache_manager import KVCacheManager
# disable yapf here as it formats differently than isort such that both fail # disable yapf here as it formats differently than isort such that both fail
# yapf: disable # yapf: disable
from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType, from vllm.v1.core.kv_cache_utils import (FreeKVCacheBlockQueue, KVCacheBlock,
FreeKVCacheBlockQueue, KVCacheBlock,
PrefixCachingMetrics, PrefixCachingMetrics,
estimate_max_model_len, estimate_max_model_len,
generate_block_hash_extra_keys, generate_block_hash_extra_keys,
@ -65,13 +65,29 @@ def new_kv_cache_spec(block_size=16,
sliding_window=sliding_window) sliding_window=sliding_window)
def test_none_hash(): def test_none_hash(monkeypatch):
assert NONE_HASH is not None import vllm.v1.core.kv_cache_utils
assert isinstance(NONE_HASH, int)
assert NONE_HASH != 0 # case 1: PYTHONHASHSEED is not set, use random
with monkeypatch.context() as m:
m.delenv('PYTHONHASHSEED', raising=False)
reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
assert reloaded_kv_cache_utils.NONE_HASH is not None
assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
assert reloaded_kv_cache_utils.NONE_HASH != 0
# case 2: PYTHONHASHSEED is set, use the seed
with monkeypatch.context() as m:
m.setenv('PYTHONHASHSEED', 'python hash seed')
reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
assert reloaded_kv_cache_utils.NONE_HASH is not None
assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
assert sha256('python hash seed') == reloaded_kv_cache_utils.NONE_HASH
def test_kv_cache_block(): def test_kv_cache_block():
import vllm.v1.core.kv_cache_utils
# Test KVCacheBlock initialization # Test KVCacheBlock initialization
block = KVCacheBlock(block_id=0) block = KVCacheBlock(block_id=0)
assert block.block_id == 0 assert block.block_id == 0
@ -85,7 +101,8 @@ def test_kv_cache_block():
assert block.ref_cnt == 0 assert block.ref_cnt == 0
# Test block hash setting and resetting # Test block hash setting and resetting
block_hash = BlockHashType(hash_value=123, token_ids=(1, 2, 3)) block_hash = vllm.v1.core.kv_cache_utils.BlockHashType(hash_value=123,
token_ids=(1, 2, 3))
block.block_hash = block_hash block.block_hash = block_hash
assert block.block_hash == block_hash assert block.block_hash == block_hash
@ -259,13 +276,14 @@ def test_generate_block_hash_extra_keys_cache_salt():
@pytest.mark.parametrize("hash_fn", [sha256, hash]) @pytest.mark.parametrize("hash_fn", [sha256, hash])
def test_hash_block_tokens(hash_fn): def test_hash_block_tokens(hash_fn):
import vllm.v1.core.kv_cache_utils
parent_block_hash = 123 parent_block_hash = 123
curr_block_token_ids = (1, 2, 3) curr_block_token_ids = (1, 2, 3)
extra_keys = ("key1", "key2") extra_keys = ("key1", "key2")
block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_hash = hash_block_tokens(hash_fn, parent_block_hash,
curr_block_token_ids, extra_keys) curr_block_token_ids, extra_keys)
assert isinstance(block_hash, BlockHashType) assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHashType)
assert block_hash.hash_value == hash_fn( assert block_hash.hash_value == hash_fn(
(parent_block_hash, curr_block_token_ids, extra_keys)) (parent_block_hash, curr_block_token_ids, extra_keys))
assert block_hash.token_ids == curr_block_token_ids assert block_hash.token_ids == curr_block_token_ids
@ -274,6 +292,7 @@ def test_hash_block_tokens(hash_fn):
@pytest.mark.parametrize("hash_fn", [sha256, hash]) @pytest.mark.parametrize("hash_fn", [sha256, hash])
def test_hash_request_tokens(hash_fn): def test_hash_request_tokens(hash_fn):
import vllm.v1.core.kv_cache_utils
request = make_request( request = make_request(
request_id=0, request_id=0,
prompt_token_ids=[_ for _ in range(6)], prompt_token_ids=[_ for _ in range(6)],
@ -288,8 +307,10 @@ def test_hash_request_tokens(hash_fn):
block_hashes = hash_request_tokens(hash_fn, block_size, request) block_hashes = hash_request_tokens(hash_fn, block_size, request)
assert len(block_hashes) == 2 assert len(block_hashes) == 2
assert isinstance(block_hashes[0], BlockHashType) assert isinstance(block_hashes[0],
assert isinstance(block_hashes[1], BlockHashType) vllm.v1.core.kv_cache_utils.BlockHashType)
assert isinstance(block_hashes[1],
vllm.v1.core.kv_cache_utils.BlockHashType)
# Check the first block # Check the first block
assert block_hashes[0].token_ids == (0, 1, 2) assert block_hashes[0].token_ids == (0, 1, 2)