mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 06:05:02 +08:00
[UT] Add ut for none hash (#17892)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
This commit is contained in:
parent
4f07a64075
commit
420caf7557
@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import importlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@ -10,8 +11,7 @@ from vllm.utils import GiB_bytes, sha256
|
||||
from vllm.v1.core.kv_cache_manager import KVCacheManager
|
||||
# disable yapf here as it formats differently than isort such that both fail
|
||||
# yapf: disable
|
||||
from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType,
|
||||
FreeKVCacheBlockQueue, KVCacheBlock,
|
||||
from vllm.v1.core.kv_cache_utils import (FreeKVCacheBlockQueue, KVCacheBlock,
|
||||
PrefixCachingMetrics,
|
||||
estimate_max_model_len,
|
||||
generate_block_hash_extra_keys,
|
||||
@ -65,13 +65,29 @@ def new_kv_cache_spec(block_size=16,
|
||||
sliding_window=sliding_window)
|
||||
|
||||
|
||||
def test_none_hash():
|
||||
assert NONE_HASH is not None
|
||||
assert isinstance(NONE_HASH, int)
|
||||
assert NONE_HASH != 0
|
||||
def test_none_hash(monkeypatch):
|
||||
import vllm.v1.core.kv_cache_utils
|
||||
|
||||
# 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():
|
||||
import vllm.v1.core.kv_cache_utils
|
||||
|
||||
# Test KVCacheBlock initialization
|
||||
block = KVCacheBlock(block_id=0)
|
||||
assert block.block_id == 0
|
||||
@ -85,7 +101,8 @@ def test_kv_cache_block():
|
||||
assert block.ref_cnt == 0
|
||||
|
||||
# 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
|
||||
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])
|
||||
def test_hash_block_tokens(hash_fn):
|
||||
import vllm.v1.core.kv_cache_utils
|
||||
parent_block_hash = 123
|
||||
curr_block_token_ids = (1, 2, 3)
|
||||
extra_keys = ("key1", "key2")
|
||||
|
||||
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
|
||||
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(
|
||||
(parent_block_hash, curr_block_token_ids, extra_keys))
|
||||
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])
|
||||
def test_hash_request_tokens(hash_fn):
|
||||
import vllm.v1.core.kv_cache_utils
|
||||
request = make_request(
|
||||
request_id=0,
|
||||
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)
|
||||
|
||||
assert len(block_hashes) == 2
|
||||
assert isinstance(block_hashes[0], BlockHashType)
|
||||
assert isinstance(block_hashes[1], BlockHashType)
|
||||
assert isinstance(block_hashes[0],
|
||||
vllm.v1.core.kv_cache_utils.BlockHashType)
|
||||
assert isinstance(block_hashes[1],
|
||||
vllm.v1.core.kv_cache_utils.BlockHashType)
|
||||
|
||||
# Check the first block
|
||||
assert block_hashes[0].token_ids == (0, 1, 2)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user