@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4f8c4b890a

Fork 0

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 14, 05:30 UTC

[Core] Use KVCacheBlock as much as possible instead of dict[block_id, KVCacheBlock] (#24830)

Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>

Jialin Ouyang committed Sep 23, 2025 at 15:11 UTC 4f8c4b890ac1ca57b2f80278641fbe3f9a6abfba

3 files changed +247 -87

tests/v1/core/test_prefix_caching.py

+127 -51

@@ -14,10 +14,11 @@ from vllm.multimodal.inputs import (MultiModalFeatureSpec,

14 MultiModalKwargsItem, PlaceholderRange)

15 from vllm.sampling_params import SamplingParams

16 from vllm.utils import sha256, sha256_cbor

17 -from vllm.v1.core.block_pool import BlockPool

17 +from vllm.v1.core.block_pool import BlockHashToBlockMap, BlockPool

18 from vllm.v1.core.kv_cache_manager import KVCacheManager, Request

19 -from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,

20 - get_block_hash, get_group_id,

19 +from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,

20 + KVCacheBlock, get_block_hash,

21 + get_group_id,

22 get_request_block_hasher,

23 hash_block_tokens, init_none_hash,

24 make_block_hash_with_group_id)

@@ -138,7 +139,7 @@ def test_prefill(hash_fn):

139 blocks = manager.allocate_slots(req0, 55,

140 len(computed_blocks.blocks[0]) * 16,

141 computed_blocks)

141 - assert blocks.get_block_ids() == ([1, 2, 3, 4], )

142 + assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4], )

143

144 # Check full block metadata

145 parent_block_hash = None

@@ -171,7 +172,7 @@ def test_prefill(hash_fn):

172 blocks = manager.allocate_slots(req1, num_new_tokens,

173 len(computed_blocks.blocks[0]) * 16,

174 computed_blocks)

174 - assert blocks.get_block_ids() == ([5], )

175 + assert blocks is not None and blocks.get_block_ids() == ([5], )

176 for block in computed_blocks.blocks[0]:

177 assert block.ref_cnt == 2

178

@@ -207,7 +208,7 @@ def test_prefill(hash_fn):

208 blocks = manager.allocate_slots(req2, num_new_tokens,

209 len(computed_blocks.blocks[0]) * 16,

210 computed_blocks)

210 - assert blocks.get_block_ids() == ([6], )

211 + assert blocks is not None and blocks.get_block_ids() == ([6], )

212

213 # Although we only have 6 free blocks, we have 8 blocks in

214 # the free block queue due to lazy removal.

@@ -227,7 +228,9 @@ def test_prefill(hash_fn):

228 len(computed_blocks.blocks[0]) * 16,

229 computed_blocks)

230 # This block ID order also checks the eviction order.

230 - assert blocks.get_block_ids() == ([7, 8, 9, 10, 4, 5, 6, 3, 2, 1], )

231 + assert blocks is not None and blocks.get_block_ids() == ([

232 + 7, 8, 9, 10, 4, 5, 6, 3, 2, 1

233 + ], )

234

235 assert free_block_queue.num_free_blocks == 0

236 assert (free_block_queue.fake_free_list_head.next_free_block

@@ -261,8 +264,9 @@ def test_prefill_hybrid_model():

264 blocks = manager.allocate_slots(req0, 55,

265 len(computed_blocks.blocks[0]) * 16,

266 computed_blocks)

264 - assert blocks.get_block_ids() == ([1, 2, 3, 4], [5, 6, 7,

265 - 8], [9, 10, 11, 12])

267 + assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4], [

268 + 5, 6, 7, 8

269 + ], [9, 10, 11, 12])

270

271 # Check full block metadata

272 parent_block_hash = None

@@ -298,7 +302,7 @@ def test_prefill_hybrid_model():

302 blocks = manager.allocate_slots(req1, num_new_tokens,

303 len(computed_blocks.blocks[0]) * 16,

304 computed_blocks)

301 - assert blocks.get_block_ids() == ([13], [14], [15])

305 + assert blocks is not None and blocks.get_block_ids() == ([13], [14], [15])

306 for block_per_group in computed_blocks.blocks:

307 for block in block_per_group:

308 if block != manager.block_pool.null_block:

@@ -309,14 +313,15 @@ def test_prefill_hybrid_model():

313 manager.free(req1)

314

315 cached_block_hash_to_block_bak = copy.copy(

312 - manager.block_pool.cached_block_hash_to_block)

316 + manager.block_pool.cached_block_hash_to_block._cache)

317

314 - def test_partial_request_hit(request_id: str, hash_to_evict: list[bytes],

318 + def test_partial_request_hit(request_id: str,

319 + hash_to_evict: list[BlockHashWithGroupId],

320 expect_hit_length: int):

321 req = make_request(request_id, common_token_ids + unique_token_ids,

322 block_size, sha256)

323 for hash_with_group_id in hash_to_evict:

319 - manager.block_pool.cached_block_hash_to_block.pop(

324 + manager.block_pool.cached_block_hash_to_block._cache.pop(

325 hash_with_group_id)

326 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)

327 assert len(req.block_hashes) == 3

@@ -324,7 +329,7 @@ def test_prefill_hybrid_model():

329 for block_per_group in computed_blocks.blocks:

330 assert len(block_per_group) == num_computed_tokens // block_size

331 for hash_with_group_id in hash_to_evict:

327 - manager.block_pool.cached_block_hash_to_block[

332 + manager.block_pool.cached_block_hash_to_block._cache[

333 hash_with_group_id] = cached_block_hash_to_block_bak[

334 hash_with_group_id]

335 manager.free(req)

@@ -362,7 +367,8 @@ def test_prefill_hybrid_model():

367 # total cache miss.

368 # The cache hit length of full attention is 1 * block_size.

369 # The cache hit length of sliding window is 2 * block_size.

365 - # Then it is cache miss as the two type of layers have different hit length.

370 + # Then it is cache miss as the two type of layers

371 + # have different hit length.

372 test_partial_request_hit("8", [

373 make_block_hash_with_group_id(block_hashes[2], 0),

374 make_block_hash_with_group_id(block_hashes[0], 1),

@@ -406,7 +412,7 @@ def test_prefill_plp():

412 blocks = manager.allocate_slots(req0, 55,

413 len(computed_blocks.blocks[0]) * 16,

414 computed_blocks)

409 - assert blocks.get_block_ids() == ([1, 2, 3, 4], )

415 + assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4], )

416 req0_block_hashes = [b.block_hash for b in blocks.blocks[0]]

417

418 # Check full block metadata

@@ -441,7 +447,7 @@ def test_prefill_plp():

447 blocks = manager.allocate_slots(req1, num_new_tokens,

448 len(computed_blocks.blocks[0]) * 16,

449 computed_blocks)

444 - assert blocks.get_block_ids() == ([5], )

450 + assert blocks is not None and blocks.get_block_ids() == ([5], )

451 for block in computed_blocks.blocks[0]:

452 assert block.ref_cnt == 2

453

@@ -478,6 +484,7 @@ def test_prefill_plp():

484 blocks = manager.allocate_slots(req2, 55,

485 len(computed_blocks.blocks[0]) * 16,

486 computed_blocks)

487 + assert blocks is not None

488 block_ids = blocks.get_block_ids()

489 # Duplicate cached blocks have different ids but same hashes vs request #0

490 assert [b.block_hash for b in blocks.blocks[0]] == req0_block_hashes

@@ -513,7 +520,7 @@ def test_decode():

520 blocks = manager.allocate_slots(req0, 55,

521 len(computed_blocks.blocks[0]) * 16,

522 computed_blocks)

516 - assert blocks.get_block_ids() == ([1, 2, 3, 4], )

523 + assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4], )

524

525 # Append slots without allocating a new block.

526 req0.num_computed_tokens = 55

@@ -558,7 +565,8 @@ def test_evict():

565 blocks = manager.allocate_slots(req0, 5 * 16 + 7,

566 len(computed_blocks.blocks[0]) * 16,

567 computed_blocks)

561 - assert len(blocks.blocks[0]) == 6 # 5 full + 1 partial

568 + # 5 full + 1 partial

569 + assert blocks is not None and len(blocks.blocks[0]) == 6

570

571 # 3 blocks.

572 req1 = make_request("1", list(range(last_token_id,

@@ -570,7 +578,7 @@ def test_evict():

578 blocks = manager.allocate_slots(req1, 3 * 16,

579 len(computed_blocks.blocks[0]) * 16,

580 computed_blocks)

573 - assert len(blocks.blocks[0]) == 3 # 3 full blocks

581 + assert blocks is not None and len(blocks.blocks[0]) == 3 # 3 full blocks

582 last_token_id += 3 * 16

583

584 # 10 - (6 + 3) == 1

@@ -592,7 +600,7 @@ def test_evict():

600 blocks = manager.allocate_slots(req2, 3,

601 len(computed_blocks.blocks[0]) * 16,

602 computed_blocks)

595 - assert blocks.get_block_ids() == ([10], )

603 + assert blocks is not None and blocks.get_block_ids() == ([10], )

604 assert manager.block_pool.free_block_queue.num_free_blocks == 7

605

606

@@ -617,7 +625,7 @@ def test_hash_block_correct_reuse():

625 blocks = manager.allocate_slots(req, num_tokens,

626 len(computed_blocks.blocks[0]) * 16,

627 computed_blocks)

620 - assert len(blocks.blocks[0]) == 1

628 + assert blocks is not None and len(blocks.blocks[0]) == 1

629

630 # Deallocate the block.

631 manager.free(req)

@@ -631,7 +639,7 @@ def test_hash_block_correct_reuse():

639 blocks = manager.allocate_slots(req, num_tokens - 1,

640 len(computed_blocks.blocks[0]) * 16,

641 computed_blocks)

634 - assert len(blocks.blocks[0]) == 1

642 + assert blocks is not None and len(blocks.blocks[0]) == 1

643

644 assert manager.block_pool.blocks[blocks.blocks[0]

645 [0].block_id].block_hash is None

@@ -658,7 +666,7 @@ def test_computed_blocks_not_evicted():

666 blocks = manager.allocate_slots(req0, num_tokens,

667 len(computed_blocks.blocks[0]) * 16,

668 computed_blocks)

661 - assert len(blocks.blocks[0]) == 1

669 + assert blocks is not None and len(blocks.blocks[0]) == 1

670 assert blocks.blocks[0][0].block_id == 1

671

672 # Allocate another block.

@@ -670,7 +678,7 @@ def test_computed_blocks_not_evicted():

678 blocks = manager.allocate_slots(req1, num_tokens,

679 len(computed_blocks.blocks[0]) * 16,

680 computed_blocks)

673 - assert len(blocks.blocks[0]) == 1

681 + assert blocks is not None and len(blocks.blocks[0]) == 1

682 assert blocks.blocks[0][0].block_id == 2

683

684 # Free the blocks.

@@ -688,7 +696,7 @@ def test_computed_blocks_not_evicted():

696 blocks = manager.allocate_slots(req2, num_tokens * 2 - num_tokens,

697 len(computed_blocks.blocks[0]) * 16,

698 computed_blocks)

691 - assert len(blocks.blocks[0]) == 1

699 + assert blocks is not None and len(blocks.blocks[0]) == 1

700 assert blocks.blocks[0][0].block_id == 2

701

702

@@ -712,7 +720,7 @@ def test_basic_prefix_caching_disabled():

720 blocks = manager.allocate_slots(req1, 10,

721 len(computed_blocks.blocks[0]) * 16,

722 computed_blocks)

715 - assert len(blocks.blocks[0]) == 3

723 + assert blocks is not None and len(blocks.blocks[0]) == 3

724

725 # Free the blocks.

726 manager.free(req1)

@@ -726,7 +734,7 @@ def test_basic_prefix_caching_disabled():

734 blocks = manager.allocate_slots(req2, 16,

735 len(computed_blocks.blocks[0]) * 16,

736 computed_blocks)

729 - assert len(blocks.blocks[0]) == 4

737 + assert blocks is not None and len(blocks.blocks[0]) == 4

738

739 # New requests should not have any blocks.

740 req3 = make_request("3", list(range(4)), block_size, sha256)

@@ -773,7 +781,8 @@ def test_cache_blocks(hash_fn):

781 assert len(block_pool.cached_block_hash_to_block) == 2

782 assert all([block.block_hash is not None for block in blocks])

783

776 - # Test that blocks that don't start from the beginning are cached correctly.

784 + # Test that blocks that don't start from the beginning are cached

785 + # correctly.

786 blocks += [KVCacheBlock(block_id=2)]

787 block_pool.cache_full_blocks(

788 request=req,

@@ -1101,7 +1110,7 @@ def test_reset_prefix_cache():

1110 all_token_ids = full_block_token_ids + unique_token_ids

1111 req0 = make_request("0", all_token_ids, block_size, sha256)

1112 blocks = manager.allocate_slots(req0, 55)

1104 - assert blocks.get_block_ids() == ([1, 2, 3, 4], )

1113 + assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4], )

1114

1115 unique_token_ids = [4] * 7

1116 all_token_ids = full_block_token_ids + unique_token_ids

@@ -1112,7 +1121,7 @@ def test_reset_prefix_cache():

1121 blocks = manager.allocate_slots(req1, 7,

1122 len(computed_blocks.blocks[0]) * 16,

1123 computed_blocks)

1115 - assert blocks.get_block_ids() == ([5], )

1124 + assert blocks is not None and blocks.get_block_ids() == ([5], )

1125

1126 # Failed to reset prefix cache because some blocks are not freed yet.

1127 assert not manager.reset_prefix_cache()

@@ -1168,49 +1177,41 @@ def test_maybe_evict_cached_block():

1177 # Manually add all blocks to cached_blocks

1178 for block, block_hash in zip(pool.blocks, block_hashes):

1179 block.block_hash = block_hash

1171 - pool.cached_block_hash_to_block[block_hash][block.block_id] = block

1180 + pool.cached_block_hash_to_block.insert(block_hash, block)

1181

1182 block0, block1, block2, block3 = pool.blocks

1174 - assert pool.cached_block_hash_to_block == {

1183 + assert pool.cached_block_hash_to_block._cache == {

1184 block_hash0: {

1185 block0.block_id: block0,

1177 - block3.block_id: block3

1178 - },

1179 - block_hash1: {

1180 - block1.block_id: block1

1186 + block3.block_id: block3,

1187 },

1182 - block_hash2: {

1183 - block2.block_id: block2

1184 - }

1188 + block_hash1: block1,

1189 + block_hash2: block2,

1190 }

1191 # Evict block1

1192 pool._maybe_evict_cached_block(block1)

1188 - assert pool.cached_block_hash_to_block == {

1193 + assert pool.cached_block_hash_to_block._cache == {

1194 block_hash0: {

1195 block0.block_id: block0,

1196 block3.block_id: block3

1197 },

1193 - block_hash2: {

1194 - block2.block_id: block2

1195 - }

1198 + block_hash2: block2,

1199 }

1200 # Evict block0: block_hash0 entry should NOT be removed, as block3

1201 # also use the same hash

1202 pool._maybe_evict_cached_block(block0)

1200 - assert pool.cached_block_hash_to_block == {

1203 + assert pool.cached_block_hash_to_block._cache == {

1204 block_hash0: {

1205 block3.block_id: block3

1206 },

1204 - block_hash2: {

1205 - block2.block_id: block2

1206 - }

1207 + block_hash2: block2,

1208 }

1209 # Evict block2

1210 pool._maybe_evict_cached_block(block2)

1210 - assert pool.cached_block_hash_to_block == {block_hash0: {3: block3}}

1211 + assert pool.cached_block_hash_to_block._cache == {block_hash0: {3: block3}}

1212 # Evict block3

1213 pool._maybe_evict_cached_block(block3)

1213 - assert pool.cached_block_hash_to_block == {}

1214 + assert pool.cached_block_hash_to_block._cache == {}

1215

1216

1217 @pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])

@@ -1374,7 +1375,7 @@ def test_eagle_with_sliding_window():

1375 # Evict the first block in the request

1376 assert manager.block_pool.get_cached_block(

1377 block_hash_first_block, kv_cache_group_ids=[0]) is not None

1377 - manager.block_pool.cached_block_hash_to_block.pop(

1378 + manager.block_pool.cached_block_hash_to_block._cache.pop(

1379 make_block_hash_with_group_id(block_hash_first_block, 0))

1380

1381 # New request

@@ -1386,3 +1387,78 @@ def test_eagle_with_sliding_window():

1387 # there will be no matched prefix.

1388 assert len(computed_blocks.blocks[0]) == 0

1389 assert num_tokens == 0

1390 +

1391 +

1392 +def test_block_lookup_cache_single_block_per_key():

1393 + cache = BlockHashToBlockMap()

1394 + key0 = BlockHashWithGroupId(b"hash0")

1395 + key1 = BlockHashWithGroupId(b"hash1")

1396 + key2 = BlockHashWithGroupId(b"hash2")

1397 + block0 = KVCacheBlock(0)

1398 + block1 = KVCacheBlock(1)

1399 +

1400 + assert cache.get_one_block(key0) is None

1401 + assert cache.get_one_block(key1) is None

1402 + assert cache.get_one_block(key2) is None

1403 + # key0 inserted

1404 + cache.insert(key0, block0)

1405 + assert cache.get_one_block(key0) is block0

1406 + assert cache.get_one_block(key1) is None

1407 + assert cache.get_one_block(key2) is None

1408 + # key1 inserted

1409 + cache.insert(key1, block1)

1410 + assert cache.get_one_block(key0) is block0

1411 + assert cache.get_one_block(key1) is block1

1412 + assert cache.get_one_block(key2) is None

1413 + # No block poped due to block_id mismatch

1414 + assert cache.pop(key0, 100) is None

1415 + assert cache.get_one_block(key0) is block0

1416 + assert cache.get_one_block(key1) is block1

1417 + assert cache.get_one_block(key2) is None

1418 + # block poped with (key0, block ID 0)

1419 + assert cache.pop(key0, 0) is block0

1420 + assert cache.get_one_block(key0) is None

1421 + assert cache.get_one_block(key1) is block1

1422 + assert cache.get_one_block(key2) is None

1423 + # No block poped due to block_id mismatch

1424 + assert cache.pop(key0, 1) is None

1425 + assert cache.get_one_block(key0) is None

1426 + assert cache.get_one_block(key1) is block1

1427 + assert cache.get_one_block(key2) is None

1428 + # block poped with (key1, block ID 1)

1429 + assert cache.pop(key1, 1) is block1

1430 + assert cache.get_one_block(key0) is None

1431 + assert cache.get_one_block(key1) is None

1432 + assert cache.get_one_block(key2) is None

1433 +

1434 +

1435 +def test_block_lookup_cache_multi_blocks_per_key():

1436 + cache = BlockHashToBlockMap()

1437 + key0 = BlockHashWithGroupId(b"hash0")

1438 + key1 = BlockHashWithGroupId(b"hash1")

1439 + block00 = KVCacheBlock(0)

1440 + block01 = KVCacheBlock(1)

1441 + block10 = KVCacheBlock(10)

1442 + block11 = KVCacheBlock(11)

1443 +

1444 + assert cache.get_one_block(key0) is None

1445 + assert cache.get_one_block(key1) is None

1446 +

1447 + cache.insert(key0, block00)

1448 + cache.insert(key0, block01)

1449 + cache.insert(key1, block10)

1450 + cache.insert(key1, block11)

1451 +

1452 + assert cache.get_one_block(key0) is block00

1453 + assert cache.pop(key0, 0) is block00

1454 + assert cache.get_one_block(key0) is block01

1455 + assert cache.pop(key0, 1) is block01

1456 + assert cache.get_one_block(key0) is None

1457 + assert cache.pop(key0, 2) is None

1458 +

1459 + assert cache.get_one_block(key1) is block10

1460 + assert cache.pop(key1, 10) is block10

1461 + assert cache.get_one_block(key1) is block11

1462 + assert cache.pop(key1, 11) is block11

1463 + assert cache.get_one_block(key1) is None

1464 + assert cache.pop(key1, 12) is None

tests/v1/core/test_single_type_kv_cache_manager.py

+8 -10

@@ -47,16 +47,15 @@ def test_chunked_local_attention_possible_cached_prefix():

47 BlockHash(str(i).encode()) for i in range(len(block_is_cached))

48 ]

50 - block_pool.cached_block_hash_to_block.clear()

50 + block_pool.cached_block_hash_to_block._cache.clear()

52 # Mock the block pool with the cached blocks

53 for i, (block_hash,

54 is_cached) in enumerate(zip(block_hash_list, block_is_cached)):

55 if is_cached:

56 - block_pool.cached_block_hash_to_block[

57 - make_block_hash_with_group_id(block_hash, 0)] = {

58 - i: block_pool.blocks[i + 10],

59 - }

56 + block_pool.cached_block_hash_to_block.insert(

57 + make_block_hash_with_group_id(block_hash, 0),

58 + block_pool.blocks[i + 10])

60 computed_blocks = manager.find_longest_cache_hit(

61 block_hashes=block_hash_list,

@@ -112,16 +111,15 @@ def test_sliding_window_possible_cached_prefix():

111 BlockHash(str(i).encode()) for i in range(len(block_is_cached))

112 ]

113

115 - block_pool.cached_block_hash_to_block.clear()

114 + block_pool.cached_block_hash_to_block._cache.clear()

115

116 # Mock the block pool with the cached blocks

117 for i, (block_hash,

118 is_cached) in enumerate(zip(block_hash_list, block_is_cached)):

119 if is_cached:

121 - block_pool.cached_block_hash_to_block[

122 - make_block_hash_with_group_id(block_hash, 0)] = {

123 - i: block_pool.blocks[i + 10],

124 - }

120 + block_pool.cached_block_hash_to_block.insert(

121 + make_block_hash_with_group_id(block_hash, 0),

122 + block_pool.blocks[i + 10])

123

124 computed_blocks = manager.find_longest_cache_hit(

125 block_hashes=block_hash_list,

vllm/v1/core/block_pool.py

+112 -26

@@ -1,8 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections import defaultdict

3 from collections.abc import Iterable

5 -from typing import Optional

4 +from typing import Any, Optional, Union

6 from vllm.distributed.kv_events import (MEDIUM_GPU, AllBlocksCleared,

7 BlockRemoved, BlockStored,

@@ -19,6 +18,103 @@ from vllm.v1.request import Request

18 logger = init_logger(__name__)

21 +class BlockHashToBlockMap:

22 + """

23 + Cache of blocks that are used for prefix caching. It caches blocks

24 + from hash directly to a block or multiple blocks

25 + (i.e. {block_hash: KVCacheBlocks})

26 + - Mostly block_hash maps to a single KVCacheBlock, and KVCacheBlocks

27 + would simply be a KVCacheBlock.

28 + - Otherwise, KVCacheBlocks is a dict from {block_id: KVCacheBlock}

29 +

30 + A cached block is a full block with a block hash that can be used

31 + for prefix caching.

32 + The cached block may be used by running requests or in the

33 + free_block_queue that could potentially be evicted.

34 +

35 + NOTE #1: We currently don't de-duplicate the blocks in the cache,

36 + meaning that if a block becomes full and is cached, we don't check

37 + if there is already an identical block in the cache. This is because

38 + we want to make sure the allocated block IDs won't change so that

39 + block tables are append-only.

40 + NOTE #2: The union type is introduced in order to reduce GC costs

41 + from the inner dict.

42 + """

43 +

44 + def __init__(self):

45 + self._cache: dict[BlockHashWithGroupId,

46 + Union[KVCacheBlock, dict[int, KVCacheBlock]]] = {}

47 +

48 + def get_one_block(self,

49 + key: BlockHashWithGroupId) -> Optional[KVCacheBlock]:

50 + """

51 + Gets any block with the given block hash key.

52 + """

53 + blocks = self._cache.get(key)

54 + if blocks is not None:

55 + if isinstance(blocks, KVCacheBlock):

56 + return blocks

57 + if isinstance(blocks, dict):

58 + return next(iter(blocks.values()))

59 + self._unexpected_blocks_type(blocks)

60 + return None

61 +

62 + def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None:

63 + """

64 + Inserts the KVCacheBlock to the cache

65 + """

66 + blocks = self._cache.get(key)

67 + if blocks is None:

68 + # When key is not found, attach a single block to the key

69 + self._cache[key] = block

70 + elif isinstance(blocks, KVCacheBlock):

71 + # If there's a block with the same key, merge the original block

72 + # and the new block into a dict

73 + self._cache[key] = {blocks.block_id: blocks, block.block_id: block}

74 + elif isinstance(blocks, dict):

75 + # If it's already a dict, simply insert the block

76 + blocks[block.block_id] = block

77 + else:

78 + self._unexpected_blocks_type(blocks)

79 +

80 + def pop(self, key: BlockHashWithGroupId,

81 + block_id: int) -> Optional[KVCacheBlock]:

82 + """

83 + Checks if block_hash exists and pop block_id from the cache

84 + """

85 + blocks = self._cache.pop(key, None)

86 + if blocks is None:

87 + # block_hash not found in the cache

88 + return None

89 + # TODO(Jialin): If key is found, block_id should always present

90 + # in blocks. We currently keep the original behaviour for safety.

91 + #

92 + # Will add block_id == blocks.block_id assertion and

93 + # use del blocks[block_id] instead as followup.

94 + if isinstance(blocks, KVCacheBlock):

95 + if blocks.block_id == block_id:

96 + return blocks

97 + # If the single block ID doesn't match, we should put the

98 + # block back (it should happen rarely)

99 + self._cache[key] = blocks

100 + return None

101 + if isinstance(blocks, dict):

102 + # Try to pop block_id from the block dict, and if dict still

103 + # contain blocks, put back to the cache.

104 + block = blocks.pop(block_id, None)

105 + if len(blocks) > 0:

106 + self._cache[key] = blocks

107 + return block

108 + self._unexpected_blocks_type(blocks)

109 + return None

110 +

111 + def __len__(self) -> int:

112 + return len(self._cache)

113 +

114 + def _unexpected_blocks_type(self, blocks: Any) -> None:

115 + raise AssertionError(f"Invalid KV cache block type {type(blocks)}")

116 +

117 +

118 class BlockPool:

119 """BlockPool that manages KVCacheBlocks.

120 It provides methods to allocate, free and cache the kv cache blocks. The

@@ -51,17 +147,9 @@ class BlockPool:

147 # enabled).

148 self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)

149

54 - # {block_hash: {block ID: block}}. A cached block is

55 - # a full block with a block hash that can be used for prefix caching.

56 - # The cached block may be used by running requests or in the

57 - # free_block_queue that could potentially be evicted.

58 - # NOTE: We currently don't de-duplicate the blocks in the cache,

59 - # meaning that if a block becomes full and is cached, we don't check

60 - # if there is already an identical block in the cache. This is because

61 - # we want to make sure the allocated block IDs won't change so that

62 - # block tables are append-only.

63 - self.cached_block_hash_to_block: dict[BlockHashWithGroupId, dict[

64 - int, KVCacheBlock]] = defaultdict(dict)

150 + # Cache for block lookup

151 + self.cached_block_hash_to_block: BlockHashToBlockMap = \

152 + BlockHashToBlockMap()

153

154 # To represent a placeholder block with block_id=0.

155 # The ref_cnt of null_block is not maintained, needs special care to

@@ -90,12 +178,11 @@ class BlockPool:

178 for group_id in kv_cache_group_ids:

179 block_hash_with_group_id = make_block_hash_with_group_id(

180 block_hash, group_id)

93 - cached_blocks_one_group = self.cached_block_hash_to_block.get(

181 + block = self.cached_block_hash_to_block.get_one_block(

182 block_hash_with_group_id)

95 - if not cached_blocks_one_group:

183 + if not block:

184 return None

97 - first_block = next(iter(cached_blocks_one_group.values()))

98 - cached_blocks.append(first_block)

185 + cached_blocks.append(block)

186 return cached_blocks

187

188 def cache_full_blocks(

@@ -140,8 +227,8 @@ class BlockPool:

227 block_hash_with_group_id = make_block_hash_with_group_id(

228 block_hash, kv_cache_group_id)

229 blk.block_hash = block_hash_with_group_id

143 - self.cached_block_hash_to_block[block_hash_with_group_id][

144 - blk.block_id] = blk

230 + self.cached_block_hash_to_block.insert(block_hash_with_group_id,

231 + blk)

232 if new_hashes is not None:

233 new_hashes.append(maybe_convert_block_hash(block_hash))

234

@@ -211,15 +298,14 @@ class BlockPool:

298 if block_hash is None:

299 # The block doesn't have hash, eviction is not needed

300 return False

214 - blocks_by_id = self.cached_block_hash_to_block.get(block_hash)

215 - if blocks_by_id is None:

216 - # block_hash not found in cached_block_hash_to_block,

301 +

302 + if self.cached_block_hash_to_block.pop(block_hash,

303 + block.block_id) is None:

304 + # block not found in cached_block_hash_to_block,

305 # eviction is not needed

306 return False

307 +

308 block.reset_hash()

220 - blocks_by_id.pop(block.block_id, None)

221 - if len(blocks_by_id) == 0:

222 - del self.cached_block_hash_to_block[block_hash]

309

310 if self.enable_kv_cache_events:

311 # FIXME (Chen): Not sure whether we should return `hash_value`

@@ -283,7 +369,7 @@ class BlockPool:

369 return False

370

371 # Remove all hashes so that no new blocks will hit.

286 - self.cached_block_hash_to_block = defaultdict(dict)

372 + self.cached_block_hash_to_block = BlockHashToBlockMap()

373

374 # Remove all hashes from all blocks.

375 for block in self.blocks: