# # 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()