582 lines
22 KiB
Python
582 lines
22 KiB
Python
#
|
|
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
|
|
import hashlib
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
import tests.ut.distributed.ascend_store._mock_deps # noqa: F401, E402
|
|
from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
|
|
AscendConnectorMetadata,
|
|
ChunkedTokenDatabase,
|
|
KeyMetadata,
|
|
LayerMultiBlockReqMeta,
|
|
LayerPoolKey,
|
|
LoadSpec,
|
|
PoolKey,
|
|
ReqMeta,
|
|
RequestTracker,
|
|
get_block_hashes,
|
|
)
|
|
|
|
_GROUPED_BLOCK_HASH_DOMAIN = b"vllm-ascend-grouped-block-hash-v1\0"
|
|
_GROUPED_BLOCK_HASH_LENGTH_PREFIX_BYTES = 4
|
|
|
|
|
|
def _expected_grouped_hash(*block_hashes):
|
|
hasher = hashlib.sha256()
|
|
hasher.update(_GROUPED_BLOCK_HASH_DOMAIN)
|
|
hasher.update(len(block_hashes).to_bytes(_GROUPED_BLOCK_HASH_LENGTH_PREFIX_BYTES, "big"))
|
|
for block_hash in block_hashes:
|
|
hash_bytes = block_hash.encode("utf-8") if isinstance(block_hash, str) else bytes(block_hash)
|
|
hasher.update(len(hash_bytes).to_bytes(_GROUPED_BLOCK_HASH_LENGTH_PREFIX_BYTES, "big"))
|
|
hasher.update(hash_bytes)
|
|
return hasher.digest()
|
|
|
|
|
|
class TestKeyMetadata(unittest.TestCase):
|
|
def test_fields(self):
|
|
meta = KeyMetadata(
|
|
model_name="llama",
|
|
head_or_tp_rank=0,
|
|
pcp_rank=0,
|
|
dcp_rank=0,
|
|
pp_rank=0,
|
|
)
|
|
self.assertEqual(meta.model_name, "llama")
|
|
self.assertEqual(meta.head_or_tp_rank, 0)
|
|
self.assertEqual(meta.pcp_rank, 0)
|
|
self.assertEqual(meta.dcp_rank, 0)
|
|
self.assertEqual(meta.pp_rank, 0)
|
|
|
|
|
|
class TestPoolKey(unittest.TestCase):
|
|
def setUp(self):
|
|
self.meta = KeyMetadata("llama", 1, 2, 3, 0)
|
|
|
|
def test_hash_equal(self):
|
|
k1 = PoolKey(self.meta, "abc123")
|
|
k2 = PoolKey(self.meta, "abc123")
|
|
self.assertEqual(hash(k1), hash(k2))
|
|
|
|
def test_hash_diff(self):
|
|
k1 = PoolKey(self.meta, "abc123")
|
|
k2 = PoolKey(self.meta, "def456")
|
|
self.assertNotEqual(hash(k1), hash(k2))
|
|
|
|
def test_to_string(self):
|
|
k = PoolKey(self.meta, "hash1")
|
|
s = k.to_string()
|
|
self.assertIn("llama", s)
|
|
self.assertIn("@pcp2", s)
|
|
self.assertIn("@dcp3", s)
|
|
self.assertIn("@head_or_tp_rank:1", s)
|
|
self.assertIn("@pp_rank:0", s)
|
|
self.assertIn("hash1", s)
|
|
|
|
def test_pp_ranks_use_distinct_keys(self):
|
|
other_pp_meta = KeyMetadata("llama", 1, 2, 3, 1)
|
|
pp0_key = PoolKey(self.meta, "hash1")
|
|
pp1_key = PoolKey(other_pp_meta, "hash1")
|
|
|
|
self.assertNotEqual(pp0_key.to_string(), pp1_key.to_string())
|
|
self.assertIn("@pp_rank:0", pp0_key.to_string())
|
|
self.assertIn("@pp_rank:1", pp1_key.to_string())
|
|
|
|
def test_split_layers(self):
|
|
k = PoolKey(self.meta, "hash1")
|
|
layers = k.split_layers(3)
|
|
self.assertEqual(len(layers), 3)
|
|
for i, lk in enumerate(layers):
|
|
self.assertIsInstance(lk, LayerPoolKey)
|
|
self.assertEqual(lk.layer_id, i)
|
|
self.assertEqual(lk.chunk_hash, "hash1")
|
|
|
|
|
|
class TestLayerPoolKey(unittest.TestCase):
|
|
def test_hash(self):
|
|
meta = KeyMetadata("model", 0, 0, 0, 0)
|
|
k1 = LayerPoolKey(meta, "h1", 0)
|
|
k2 = LayerPoolKey(meta, "h1", 1)
|
|
self.assertNotEqual(hash(k1), hash(k2))
|
|
|
|
def test_to_string_contains_layer_id(self):
|
|
meta = KeyMetadata("model", 0, 0, 0, 0)
|
|
k = LayerPoolKey(meta, "h1", 5)
|
|
s = k.to_string()
|
|
self.assertIn("@layer_id:5", s)
|
|
self.assertIn("model", s)
|
|
self.assertTrue(s.endswith("@h1"))
|
|
|
|
|
|
class TestChunkedTokenDatabase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.meta = KeyMetadata("llama", 0, 0, 0, 0)
|
|
self.db = ChunkedTokenDatabase([self.meta], block_size=[16], partitions=None)
|
|
self.db.set_group_buffers({0: [1000, 2000]}, {0: [160, 320]}, group_num_layers={0: 1})
|
|
|
|
def test_make_key_by_hash(self):
|
|
key = self.db._make_key_by_hash("abc")
|
|
self.assertIsInstance(key, PoolKey)
|
|
self.assertEqual(key.chunk_hash, "abc")
|
|
|
|
def test_process_tokens_empty(self):
|
|
result = list(self.db.process_tokens(32, []))
|
|
self.assertEqual(result, [])
|
|
|
|
def test_process_tokens_with_str_hashes(self):
|
|
hashes = ["aaa", "bbb"]
|
|
result = list(self.db.process_tokens(32, hashes))
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0][0], 0) # start
|
|
self.assertEqual(result[0][1], 16) # end
|
|
self.assertEqual(result[1][0], 16)
|
|
self.assertEqual(result[1][1], 32)
|
|
|
|
def test_process_tokens_with_bytes_hashes(self):
|
|
hashes = [b"\xaa\xbb", b"\xcc\xdd"]
|
|
result = list(self.db.process_tokens(32, hashes))
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_process_tokens_with_mask(self):
|
|
hashes = ["a", "b", "c"]
|
|
result = list(self.db.process_tokens(48, hashes, mask_num=16))
|
|
# first chunk (start=0 < mask_num=16) should be skipped
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0][0], 16)
|
|
|
|
def test_process_tokens_with_tail_clipped_block_ids_maps_tail_chunks(self):
|
|
db = ChunkedTokenDatabase([self.meta], block_size=[128], partitions=None)
|
|
hashes = [bytes([idx % 251]) * 32 for idx in range(128)]
|
|
|
|
result = list(
|
|
db.process_token_key_strings_with_block_ids(
|
|
128 * 128,
|
|
hashes,
|
|
[1000, 1001, 1002, 1003],
|
|
)
|
|
)
|
|
|
|
self.assertEqual(
|
|
[start for start, _, _, _, _ in result],
|
|
[124 * 128, 125 * 128, 126 * 128, 127 * 128],
|
|
)
|
|
self.assertEqual(
|
|
[block_id for _, _, _, _, block_id in result],
|
|
[1000, 1001, 1002, 1003],
|
|
)
|
|
|
|
def test_process_tokens_token_len_shorter_than_all_blocks(self):
|
|
hashes = ["a", "b", "c", "d"]
|
|
# token_len=32 means only first 2 blocks valid
|
|
result = list(self.db.process_tokens(32, hashes))
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_process_tokens_rehashes_grouped_hashes(self):
|
|
db = ChunkedTokenDatabase([self.meta], block_size=[16], partitions=None, hash_block_size=8)
|
|
result = list(db.process_tokens(32, ["a", "b", "c", "d"]))
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0][2].chunk_hash, _expected_grouped_hash("a", "b").hex())
|
|
self.assertEqual(len(result[0][2].chunk_hash), 64)
|
|
|
|
def test_key_strings_match_pool_keys(self):
|
|
hashes = ["aaa", "bbb", "ccc"]
|
|
pool_keys = list(self.db.process_tokens(40, hashes))
|
|
self.assertEqual(
|
|
list(self.db.process_token_key_strings(40, hashes)),
|
|
[
|
|
(start, end, key.to_string(), hash_val)
|
|
for (start, end, key), hash_val in zip(pool_keys, hashes, strict=True)
|
|
],
|
|
)
|
|
|
|
block_ids = [5, 6]
|
|
self.assertEqual(
|
|
list(self.db.process_token_key_strings_with_block_ids(32, hashes, block_ids)),
|
|
[
|
|
(start, end, key.to_string(), hash_val, block_id)
|
|
for (start, end, key), hash_val, block_id in zip(pool_keys[:2], hashes[:2], block_ids, strict=True)
|
|
],
|
|
)
|
|
|
|
def test_direct_keys_preserve_multigroup_layerwise_key_semantics(self):
|
|
group_metadata = [
|
|
KeyMetadata("llama", 0, 0, 0, 0),
|
|
KeyMetadata("llama", 1, 0, 0, 0),
|
|
]
|
|
db = ChunkedTokenDatabase(group_metadata, block_size=[16, 32], partitions=None, hash_block_size=16)
|
|
db.set_group_buffers(
|
|
{0: [1000], 1: [2000]},
|
|
{0: [160], 1: [320]},
|
|
group_cache_families={0: "c1", 1: "c2"},
|
|
group_num_layers={0: 2, 1: 2},
|
|
)
|
|
hashes = ["a", "b", "c", "d"]
|
|
|
|
pool_key_result = list(db.process_tokens(64, hashes, kv_cache_group_id=1))
|
|
direct_key_result = list(db.process_token_key_strings(64, hashes, kv_cache_group_id=1))
|
|
|
|
self.assertEqual(len(pool_key_result), 1)
|
|
self.assertEqual(
|
|
direct_key_result[0][:3],
|
|
(pool_key_result[0][0], pool_key_result[0][1], pool_key_result[0][2].to_string()),
|
|
)
|
|
layer_key = pool_key_result[0][2].split_layers(2)[1]
|
|
self.assertIn("@group:1@cache_role:kv@cache_family:c2@layer_id:1", layer_key.to_string())
|
|
|
|
def test_key_strings_pre_shard_after_filtering(self):
|
|
hashes = ["a", "b", "c", "d"]
|
|
store_mask = [True, False, True, True]
|
|
result = list(
|
|
self.db.process_token_key_strings_with_block_ids(
|
|
64,
|
|
hashes,
|
|
[10, 11, 12, 13],
|
|
chunk_filter=lambda start: store_mask[start // 16],
|
|
shard_rank=1,
|
|
shard_size=2,
|
|
)
|
|
)
|
|
self.assertEqual([(start, end, block_id) for start, end, _, _, block_id in result], [(32, 48, 12)])
|
|
|
|
def test_get_block_hashes_rehashes_groups(self):
|
|
for hashes in (["a", "b", "c", "d"], [b"a", b"b", b"c", b"d"]):
|
|
with self.subTest(hash_type=type(hashes[0])):
|
|
result = get_block_hashes(hashes, group_block_size=32, hash_block_size=16)
|
|
expected = [
|
|
_expected_grouped_hash(hashes[0], hashes[1]),
|
|
_expected_grouped_hash(hashes[2], hashes[3]),
|
|
]
|
|
self.assertEqual(list(result), expected)
|
|
|
|
def test_prepare_value(self):
|
|
addr, size, block_id = self.db.prepare_value(0, 16, [5, 6, 7])
|
|
self.assertEqual(block_id, 5)
|
|
self.assertEqual(len(addr), 2)
|
|
self.assertEqual(addr[0], 1000 + 5 * 160)
|
|
self.assertEqual(addr[1], 2000 + 5 * 320)
|
|
self.assertEqual(size[0], 160)
|
|
self.assertEqual(size[1], 320)
|
|
|
|
def test_prepare_value_partial_block(self):
|
|
addr, size, block_id = self.db.prepare_value(0, 8, [5])
|
|
self.assertEqual(size[0], 80) # 160/16*8
|
|
self.assertEqual(size[1], 160) # 320/16*8
|
|
|
|
def test_prepare_value_uses_block_id_override(self):
|
|
addr, size, block_id = self.db.prepare_value(64, 80, [5], block_id=99)
|
|
self.assertEqual(block_id, 99)
|
|
self.assertEqual(addr[0], 1000 + 99 * 160)
|
|
self.assertEqual(addr[1], 2000 + 99 * 320)
|
|
self.assertEqual(size[0], 160)
|
|
self.assertEqual(size[1], 320)
|
|
|
|
def test_prepare_value_layer(self):
|
|
addr, size, block_id = self.db.prepare_value_layer(0, 16, [5, 6], layer_id=0)
|
|
self.assertEqual(block_id, 5)
|
|
self.assertEqual(len(addr), 2)
|
|
# layer_id=0, entries_per_layers=2 => group_addrs[0] and group_addrs[1]
|
|
self.assertEqual(addr[0], 1000 + 5 * 160)
|
|
self.assertEqual(addr[1], 2000 + 5 * 320)
|
|
|
|
def test_decode_adaptor_prefill_pp_no_partitions(self):
|
|
key, addr, size = self.db.decode_adaptor_prefill_pp(["k1"], [[1, 2]], [[10, 20]])
|
|
self.assertEqual(key, ["k1"])
|
|
|
|
def test_decode_adaptor_prefill_pp_single_partition(self):
|
|
db = ChunkedTokenDatabase([self.meta], [16], partitions=[4])
|
|
key, addr, size = db.decode_adaptor_prefill_pp(["k1"], [[1, 2]], [[10, 20]])
|
|
self.assertEqual(key, ["k1"])
|
|
|
|
def test_decode_adaptor_prefill_pp_multi_partition(self):
|
|
db = ChunkedTokenDatabase([self.meta], [16], partitions=[2, 2])
|
|
db.set_group_buffers({0: [1000, 2000]}, {0: [160, 320]})
|
|
keys = ["k1@pp_rank:0"]
|
|
addrs = [[1, 2, 3, 4, 5, 6, 7, 8]]
|
|
sizes = [[10, 20, 30, 40, 50, 60, 70, 80]]
|
|
new_keys, new_addrs, new_sizes = db.decode_adaptor_prefill_pp(keys, addrs, sizes)
|
|
self.assertEqual(len(new_keys), 2)
|
|
self.assertIn("@pp_rank:0", new_keys[0])
|
|
self.assertIn("@pp_rank:1", new_keys[1])
|
|
|
|
|
|
class TestLoadSpec(unittest.TestCase):
|
|
def test_fields(self):
|
|
spec = LoadSpec(vllm_cached_tokens=10, kvpool_cached_tokens=20, can_load=True)
|
|
self.assertEqual(spec.vllm_cached_tokens, 10)
|
|
self.assertEqual(spec.kvpool_cached_tokens, 20)
|
|
self.assertTrue(spec.can_load)
|
|
self.assertEqual(spec.token_len, 0)
|
|
|
|
def test_token_len_default(self):
|
|
spec = LoadSpec(0, 0, False, token_len=128)
|
|
self.assertEqual(spec.token_len, 128)
|
|
|
|
|
|
class TestRequestTracker(unittest.TestCase):
|
|
def test_from_new_request(self):
|
|
new_req = MagicMock()
|
|
new_req.req_id = "req-1"
|
|
new_req.block_ids = [10, 20, 30]
|
|
new_req.prompt_token_ids = list(range(100))
|
|
|
|
tracker = RequestTracker.from_new_request(new_req, num_tokens_to_compute=48)
|
|
self.assertEqual(tracker.req_id, "req-1")
|
|
self.assertEqual(tracker.token_len, 48)
|
|
self.assertEqual(tracker.allocated_block_ids, [10, 20, 30])
|
|
self.assertEqual(len(tracker.token_ids), 48)
|
|
self.assertEqual(tracker.num_saved_tokens, 0)
|
|
|
|
def test_from_new_request_nested_block_ids(self):
|
|
new_req = MagicMock()
|
|
new_req.req_id = "req-2"
|
|
new_req.block_ids = [[10, 20], [30, 40]]
|
|
new_req.prompt_token_ids = list(range(32))
|
|
|
|
tracker = RequestTracker.from_new_request(new_req, num_tokens_to_compute=32)
|
|
self.assertEqual(tracker.allocated_block_ids, [10, 20])
|
|
|
|
def test_update_with_list(self):
|
|
tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1, 2])
|
|
tracker.update([3, 4])
|
|
self.assertEqual(tracker.allocated_block_ids, [1, 2, 3, 4])
|
|
|
|
def test_update_with_tuple(self):
|
|
tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1])
|
|
tracker.update(([5, 6], [7, 8]))
|
|
self.assertEqual(tracker.allocated_block_ids, [1, 5, 6])
|
|
|
|
def test_update_with_empty(self):
|
|
tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1])
|
|
tracker.update([])
|
|
self.assertEqual(tracker.allocated_block_ids, [1])
|
|
|
|
def test_update_invalid_type(self):
|
|
tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1])
|
|
with self.assertRaises(ValueError):
|
|
tracker.update("invalid") # type: ignore[arg-type]
|
|
|
|
def test_update_mamba_with_tuple(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1", token_len=16, allocated_block_ids_by_group=[[1], [2], [3], [4]], block_sizes=[16] * 4
|
|
)
|
|
tracker.update(([5, 6], [0, 7], [0, 8], [0, 9]))
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[0], [1, 5, 6])
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[1], [2, 0, 7])
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[2], [3, 0, 8])
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[3], [4, 0, 9])
|
|
|
|
def test_update_mamba_mtp_with_tuple_chunk2(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids_by_group=[
|
|
[1, 2],
|
|
[0, 3, 4, 5, 6],
|
|
[0, 7, 8, 9, 10],
|
|
[0, 11, 12, 13, 14],
|
|
],
|
|
mamba_group_ids=[1, 2, 3],
|
|
num_speculative_blocks=3,
|
|
block_sizes=[16] * 4,
|
|
)
|
|
|
|
tracker.update(([15, 16], [4, 17], [8, 18], [12, 19]), 32)
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[0], [1, 2, 15, 16])
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[1], [0, 3, 0, 5, 6, 4, 17])
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[2], [0, 7, 0, 9, 10, 8, 18])
|
|
self.assertEqual(tracker.allocated_block_ids_by_group[3], [0, 11, 0, 13, 14, 12, 19])
|
|
|
|
def test_update_mamba_mtp_with_tuple_chunk8(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=128,
|
|
allocated_block_ids_by_group=[
|
|
[1, 2, 3, 4, 5, 6, 7, 8],
|
|
[0, 0, 0, 0, 0, 0, 0, 9, 10, 11, 12],
|
|
[0, 0, 0, 0, 0, 0, 0, 13, 14, 15, 16],
|
|
[0, 0, 0, 0, 0, 0, 0, 17, 18, 19, 20],
|
|
],
|
|
mamba_group_ids=[1, 2, 3],
|
|
num_speculative_blocks=3,
|
|
block_sizes=[16] * 4,
|
|
)
|
|
|
|
tracker.update(
|
|
(
|
|
[21, 22, 23, 24, 25, 26, 27, 28],
|
|
[0, 0, 0, 0, 10, 11, 12, 29],
|
|
[0, 0, 0, 0, 14, 15, 16, 30],
|
|
[0, 0, 0, 0, 18, 19, 20, 31],
|
|
),
|
|
128,
|
|
)
|
|
self.assertEqual(
|
|
tracker.allocated_block_ids_by_group[0], [1, 2, 3, 4, 5, 6, 7, 8, 21, 22, 23, 24, 25, 26, 27, 28]
|
|
)
|
|
self.assertEqual(
|
|
tracker.allocated_block_ids_by_group[1], [0, 0, 0, 0, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0, 10, 11, 12, 29]
|
|
)
|
|
self.assertEqual(
|
|
tracker.allocated_block_ids_by_group[2], [0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 15, 16, 30]
|
|
)
|
|
self.assertEqual(
|
|
tracker.allocated_block_ids_by_group[3], [0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 19, 20, 31]
|
|
)
|
|
|
|
|
|
class TestReqMeta(unittest.TestCase):
|
|
def test_from_request_tracker_basic_save(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=0,
|
|
token_ids=list(range(32)),
|
|
)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, block_hashes=[b"h1", b"h2"])
|
|
self.assertIsNotNone(meta)
|
|
self.assertEqual(meta.req_id, "r1")
|
|
self.assertTrue(meta.can_save)
|
|
self.assertEqual(meta.token_len_chunk, 32)
|
|
self.assertIsNone(meta.load_spec)
|
|
|
|
def test_from_request_tracker_skip_save(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=0,
|
|
)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, skip_save=True)
|
|
self.assertIsNone(meta)
|
|
|
|
def test_from_request_tracker_with_load_spec(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=0,
|
|
)
|
|
load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=32, can_load=True)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, load_spec=load_spec, skip_save=True)
|
|
self.assertIsNotNone(meta)
|
|
self.assertIsNotNone(meta.load_spec)
|
|
|
|
def test_from_request_tracker_load_spec_cannot_load(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=32,
|
|
)
|
|
load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=32, can_load=False)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, load_spec=load_spec, skip_save=True)
|
|
# can_load=False => load_spec set to None in meta,
|
|
# but skip_save+load_spec input is not None, so meta is still created
|
|
self.assertIsNotNone(meta)
|
|
self.assertIsNone(meta.load_spec)
|
|
self.assertFalse(meta.can_save)
|
|
|
|
def test_from_request_tracker_partial_tokens_discarded(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=20,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=0,
|
|
)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, discard_partial_chunks=True)
|
|
self.assertIsNotNone(meta)
|
|
self.assertEqual(meta.token_len_chunk, 16)
|
|
|
|
def test_from_request_tracker_no_discard(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=20,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=0,
|
|
)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, discard_partial_chunks=False)
|
|
self.assertIsNotNone(meta)
|
|
self.assertEqual(meta.token_len_chunk, 20)
|
|
|
|
def test_from_request_tracker_already_saved(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=32,
|
|
)
|
|
meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16)
|
|
|
|
# num_saved_tokens=32, chunk_boundary=ceil(33/16)*16=48 > 32
|
|
# so skip_save, and no load_spec => None
|
|
self.assertIsNone(meta)
|
|
|
|
def test_from_request_tracker_with_original_block_size(self):
|
|
tracker = RequestTracker(
|
|
req_id="r1",
|
|
token_len=32,
|
|
allocated_block_ids=[0, 1],
|
|
num_saved_tokens=0,
|
|
)
|
|
# Provide block_hashes (2 full blocks for token_len=32 / granularity=16)
|
|
# so the boundary_without_hash short-circuit does not zero out the save
|
|
# length and skip; this exercises the original_block_size propagation.
|
|
meta = ReqMeta.from_request_tracker(
|
|
tracker,
|
|
cache_transfer_granularity=16,
|
|
original_block_size=8,
|
|
block_hashes=[b"h0", b"h1"],
|
|
)
|
|
self.assertIsNotNone(meta)
|
|
self.assertEqual(meta.original_block_size, 8)
|
|
|
|
|
|
class TestAscendConnectorMetadata(unittest.TestCase):
|
|
def test_add_request(self):
|
|
meta = AscendConnectorMetadata(unfinished_request_ids=set(), preempted_req_ids=set())
|
|
req = ReqMeta(
|
|
req_id="r1",
|
|
token_len_chunk=16,
|
|
block_ids=[0],
|
|
block_hashes=[],
|
|
)
|
|
meta.add_request(req)
|
|
self.assertEqual(len(meta.requests), 1)
|
|
self.assertEqual(meta.requests[0].req_id, "r1")
|
|
|
|
|
|
class TestLayerMultiBlockReqMeta(unittest.TestCase):
|
|
def test_fields(self):
|
|
meta = LayerMultiBlockReqMeta(
|
|
req_id="r1",
|
|
keys=[],
|
|
starts=[0, 16],
|
|
ends=[16, 32],
|
|
block_ids=[0, 1],
|
|
layer_id=2,
|
|
)
|
|
self.assertEqual(meta.req_id, "r1")
|
|
self.assertEqual(meta.layer_id, 2)
|
|
self.assertTrue(meta.is_last_chunk)
|
|
self.assertIsNone(meta.current_event)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|