@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 48ddb02b79

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 12, 13:30 UTC

[Hybrid Allocator] Support KV cache groups with different block_size (#29143)

Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com>

Yifan Qiao committed Nov 25, 2025 at 07:30 UTC 48ddb02b79d7e22e2eefbf5294bf70de50afd1b2

11 files changed +472 -87

tests/v1/core/test_kv_cache_utils.py

+45 -4

@@ -1248,7 +1248,9 @@ def test_allocate_with_lookahead():

1248 )

1249

1250 # Test case 1: Requires additional lookahead tokens

1251 - kv_cache_manager = KVCacheManager(kv_cache_config=config, max_model_len=100)

1251 + kv_cache_manager = KVCacheManager(

1252 + kv_cache_config=config, max_model_len=100, hash_block_size=block_size

1253 + )

1254 blocks = kv_cache_manager.allocate_slots(

1255 request,

1256 num_new_tokens=3,

@@ -1257,7 +1259,9 @@ def test_allocate_with_lookahead():

1259 assert len(blocks.get_block_ids()[0]) == 2 # ceil(5/4)=2 blocks

1260

1261 # Test case 2: With precomputed blocks

1260 - kv_cache_manager = KVCacheManager(kv_cache_config=config, max_model_len=100)

1262 + kv_cache_manager = KVCacheManager(

1263 + kv_cache_config=config, max_model_len=100, hash_block_size=block_size

1264 + )

1265 # required_blocks = ceil((3 + 2) /4) = 2

1266 blocks = kv_cache_manager.allocate_slots(

1267 request,

@@ -1268,7 +1272,9 @@ def test_allocate_with_lookahead():

1272

1273 # Test case 3: With precomputed blocks

1274 # required_blocks = ceil((3 + 4) / 4) = 2

1271 - kv_cache_manager = KVCacheManager(kv_cache_config=config, max_model_len=100)

1275 + kv_cache_manager = KVCacheManager(

1276 + kv_cache_config=config, max_model_len=100, hash_block_size=block_size

1277 + )

1278 blocks = kv_cache_manager.allocate_slots(

1279 request,

1280 num_new_tokens=3,

@@ -1495,7 +1501,8 @@ def test_get_kv_cache_config_one_worker():

1501 ),

1502 ],

1503 )

1498 - # different hidden size

1504 +

1505 + # different hidden size but same type, use UniformTypeKVCacheSpecs

1506 kv_cache_specs_hybrid = {

1507 "layer_1": new_kv_cache_spec(head_size=128),

1508 "layer_2": new_kv_cache_spec(head_size=64),

@@ -1519,6 +1526,40 @@ def test_get_kv_cache_config_one_worker():

1526 ],

1527 )

1528

1529 + # Different hidden size and different type, align by different block size

1530 + kv_cache_specs_hybrid = {

1531 + "layer_1": new_kv_cache_spec(head_size=64),

1532 + "layer_2": new_sliding_window_spec(head_size=32),

1533 + }

1534 + kv_cache_config_hybrid = get_kv_cache_configs(

1535 + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32]

1536 + )[0]

1537 + assert kv_cache_config_hybrid == KVCacheConfig(

1538 + num_blocks=32,

1539 + kv_cache_tensors=[

1540 + KVCacheTensor(

1541 + size=mem_per_block_per_layer * 32, shared_by=["layer_1", "layer_2"]

1542 + ),

1543 + ],

1544 + kv_cache_groups=[

1545 + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)),

1546 + KVCacheGroupSpec(

1547 + ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32)

1548 + ),

1549 + ],

1550 + )

1551 +

1552 + # different hidden size that cannot be aligned by using different block size

1553 + kv_cache_specs_hybrid = {

1554 + "layer_1": new_kv_cache_spec(head_size=64),

1555 + "layer_2": new_sliding_window_spec(head_size=96),

1556 + }

1557 +

1558 + with pytest.raises(NotImplementedError):

1559 + get_kv_cache_configs(

1560 + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]

1561 + )[0]

1562 +

1563 # Test num_gpu_blocks_override

1564 vllm_config.cache_config.num_gpu_blocks_override = 16

1565 kv_cache_config_override_blocks = get_kv_cache_configs(

tests/v1/core/test_prefix_caching.py

+93 -2

@@ -134,6 +134,7 @@ def test_prefill(hash_fn):

134 make_kv_cache_config(block_size, 11),

135 max_model_len=8192,

136 enable_caching=True,

137 + hash_block_size=block_size,

138 )

139

140 # Complete 3 blocks (48 tokens)

@@ -256,6 +257,7 @@ def test_prefill_hybrid_model():

257 make_kv_cache_config_hybrid_model(block_size, 21),

258 max_model_len=8192,

259 enable_caching=True,

260 + hash_block_size=block_size,

261 )

262

263 hash_fn = sha256

@@ -416,6 +418,7 @@ def test_prefill_plp():

418 make_kv_cache_config(block_size, 11),

419 max_model_len=8192,

420 enable_caching=True,

421 + hash_block_size=block_size,

422 )

423 # the default hash function is sha256

424 hash_fn = sha256

@@ -523,6 +526,7 @@ def test_decode():

526 make_kv_cache_config(block_size, 11),

527 max_model_len=8192,

528 enable_caching=True,

529 + hash_block_size=block_size,

530 )

531

532 # Complete 3 blocks (48 tokens)

@@ -585,6 +589,7 @@ def test_evict():

589 make_kv_cache_config(block_size, 11),

590 max_model_len=8192,

591 enable_caching=True,

592 + hash_block_size=block_size,

593 )

594

595 last_token_id = 5 * 16 + 7

@@ -643,6 +648,7 @@ def test_hash_block_correct_reuse():

648 make_kv_cache_config(16, 2),

649 max_model_len=8192,

650 enable_caching=True,

651 + hash_block_size=block_size,

652 )

653

654 # Allocate 1 block and cache it.

@@ -683,6 +689,7 @@ def test_computed_blocks_not_evicted():

689 make_kv_cache_config(block_size, 3),

690 max_model_len=8192,

691 enable_caching=True,

692 + hash_block_size=block_size,

693 )

694

695 # Allocate a block and cache it.

@@ -741,6 +748,7 @@ def test_basic_prefix_caching_disabled():

748 make_kv_cache_config(block_size, 5),

749 max_model_len=8192,

750 enable_caching=False,

751 + hash_block_size=block_size,

752 )

753

754 req1 = make_request(

@@ -790,6 +798,7 @@ def test_cache_blocks(hash_fn):

798 block_pool = BlockPool(

799 num_gpu_blocks=5,

800 enable_caching=True,

801 + hash_block_size=block_size,

802 )

803 # Req:

804 # Block 0: [0, 1, 2, 3]

@@ -833,7 +842,9 @@ def test_cache_blocks_multi_group():

842 This tests that blocks are cached correctly for different kv cache groups.

843 """

844 block_size = 4

836 - block_pool = BlockPool(num_gpu_blocks=10, enable_caching=True)

845 + block_pool = BlockPool(

846 + num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size

847 + )

848

849 # Req:

850 # Block 0/4: [0, 1, 2, 3]

@@ -921,6 +932,7 @@ def test_mm_prefix_caching():

932 make_kv_cache_config(block_size, 11),

933 max_model_len=8192,

934 enable_caching=True,

935 + hash_block_size=block_size,

936 )

937

938 # Common prompt tokens (T is text tokens and P is image placeholder tokens)

@@ -1020,6 +1032,7 @@ def test_cache_key_salting():

1032 make_kv_cache_config(block_size, 11),

1033 max_model_len=8192,

1034 enable_caching=True,

1035 + hash_block_size=block_size,

1036 )

1037

1038 # 3 complete blocks and an incomplete block with 11 tokens.

@@ -1101,6 +1114,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():

1114 make_kv_cache_config(block_size, 11),

1115 max_model_len=8192,

1116 enable_caching=True,

1117 + hash_block_size=block_size,

1118 )

1119 # Complete 3 blocks (48 tokens)

@@ -1173,6 +1187,7 @@ def test_reset_prefix_cache():

1187 make_kv_cache_config(block_size, 11),

1188 max_model_len=8192,

1189 enable_caching=True,

1190 + hash_block_size=block_size,

1191 )

1192

1193 full_block_token_ids = [i for i in range(3) for _ in range(16)]

@@ -1213,6 +1228,7 @@ def test_prefix_cache_stats_disabled():

1228 make_kv_cache_config(block_size, 11),

1229 max_model_len=8192,

1230 enable_caching=True,

1231 + hash_block_size=block_size,

1232 log_stats=False, # Disable logging stats

1233 )

1234 assert manager.prefix_cache_stats is None

@@ -1232,7 +1248,7 @@ def test_prefix_cache_stats_disabled():

1248

1249

1250 def test_maybe_evict_cached_block():

1235 - pool = BlockPool(num_gpu_blocks=4, enable_caching=True)

1251 + pool = BlockPool(num_gpu_blocks=4, enable_caching=True, hash_block_size=16)

1252 block_hash0 = make_block_hash_with_group_id(BlockHash(b"10"), 1000)

1253 block_hash1 = make_block_hash_with_group_id(BlockHash(b"20"), 2000)

1254 block_hash2 = make_block_hash_with_group_id(BlockHash(b"30"), 3000)

@@ -1293,6 +1309,7 @@ def test_kv_cache_events(blocks_to_cache: int):

1309 max_model_len=8192,

1310 enable_caching=True,

1311 enable_kv_cache_events=True,

1312 + hash_block_size=block_size,

1313 )

1314

1315 num_tokens = block_size * blocks_to_cache

@@ -1351,6 +1368,7 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int):

1368 max_model_len=8192,

1369 enable_caching=True,

1370 enable_kv_cache_events=True,

1371 + hash_block_size=block_size,

1372 )

1373

1374 # Test with LoRA request

@@ -1405,6 +1423,7 @@ def test_eagle_enabled_removes_last_block():

1423 max_model_len=8192,

1424 enable_caching=True,

1425 use_eagle=True,

1426 + hash_block_size=block_size,

1427 )

1428

1429 # Request with 3 full blocks (48 tokens)

@@ -1437,6 +1456,7 @@ def test_eagle_with_partial_blocks():

1456 max_model_len=8192,

1457 enable_caching=True,

1458 use_eagle=True,

1459 + hash_block_size=block_size,

1460 )

1461 # 2 full blocks + 5 tokens (non-divisible length)

1462 token_ids = [0] * (2 * block_size + 5)

@@ -1476,6 +1496,7 @@ def test_eagle_with_sliding_window():

1496 max_model_len=8192,

1497 enable_caching=True,

1498 use_eagle=True,

1499 + hash_block_size=block_size,

1500 )

1501

1502 # 2 full blocks + 5 tokens (non-divisible length)

@@ -1522,6 +1543,76 @@ def test_eagle_with_sliding_window():

1543 assert num_tokens == 0

1544

1545

1546 +def test_different_block_size():

1547 + block_size = 16

1548 + # full attention and sliding window attention layers have the same page size:

1549 + # (32 tokens/block * float16 token, vs. 16 tokens/block * float32 token)

1550 + kv_cache_config = KVCacheConfig(

1551 + num_blocks=100,

1552 + kv_cache_tensors=[],

1553 + kv_cache_groups=[

1554 + KVCacheGroupSpec(

1555 + ["layer1"],

1556 + FullAttentionSpec(block_size * 2, 1, 1, torch.float16),

1557 + ),

1558 + KVCacheGroupSpec(

1559 + ["layer2"],

1560 + SlidingWindowSpec(

1561 + block_size,

1562 + 1,

1563 + 1,

1564 + torch.float32,

1565 + sliding_window=2 * block_size,

1566 + ),

1567 + ),

1568 + ],

1569 + )

1570 + manager = KVCacheManager(

1571 + kv_cache_config=kv_cache_config,

1572 + max_model_len=8192,

1573 + enable_caching=True,

1574 + hash_block_size=block_size,

1575 + )

1576 +

1577 + # 10 blocks of 16 tokens each. Token ids are not strictly aligned for each block.

1578 + common_token_ids = [i for i in range(10) for _ in range(block_size)]

1579 +

1580 + req0 = make_request("0", common_token_ids, block_size, sha256)

1581 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

1582 + assert not computed_blocks.blocks[0]

1583 + assert not computed_blocks.blocks[1]

1584 + assert num_computed_tokens == 0

1585 + blocks = manager.allocate_slots(

1586 + req0, 7 * block_size, len(computed_blocks.blocks[0]) * 16, computed_blocks

1587 + )

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

1589 + req1 = make_request("1", common_token_ids[: 7 * block_size + 1], block_size, sha256)

1590 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1591 + assert len(computed_blocks.blocks[0]) == 3

1592 + assert len(computed_blocks.blocks[1]) == 6

1593 + assert num_computed_tokens == 6 * 16

1594 +

1595 + req2 = make_request("2", common_token_ids[: 6 * block_size + 1], block_size, sha256)

1596 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

1597 + assert len(computed_blocks.blocks[0]) == 3

1598 + assert len(computed_blocks.blocks[1]) == 6

1599 + assert num_computed_tokens == 6 * 16

1600 +

1601 + # Evict some blocks to make sliding window cache hit length 5*16

1602 + # But should return 4 * 16 because full attention cache hit length must be

1603 + # a multiple of 32

1604 + manager.block_pool.cached_block_hash_to_block.pop(

1605 + make_block_hash_with_group_id(req1.block_hashes[6], 1), 11

1606 + )

1607 + manager.block_pool.cached_block_hash_to_block.pop(

1608 + make_block_hash_with_group_id(req1.block_hashes[5], 1), 10

1609 + )

1610 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

1611 + assert len(computed_blocks.blocks[0]) == 2

1612 + assert len(computed_blocks.blocks[1]) == 4

1613 + assert num_computed_tokens == 4 * 16

1614 +

1615 +

1616 def test_block_lookup_cache_single_block_per_key():

1617 cache = BlockHashToBlockMap()

1618 key0 = BlockHashWithGroupId(b"hash0")

tests/v1/core/test_single_type_kv_cache_manager.py

+16 -6

@@ -41,7 +41,9 @@ def test_chunked_local_attention_possible_cached_prefix():

41 attention_chunk_size=4,

42 )

44 - block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)

44 + block_pool = BlockPool(

45 + num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size

46 + )

47 manager = get_chunked_local_attention_manager(

48 chunked_local_attention_spec, block_pool

49 )

@@ -70,6 +72,7 @@ def test_chunked_local_attention_possible_cached_prefix():

72 block_pool=block_pool,

73 kv_cache_spec=chunked_local_attention_spec,

74 use_eagle=False,

75 + alignment_tokens=block_size,

76 )[0]

77 assert len(computed_blocks) == expect_length

@@ -111,7 +114,9 @@ def test_sliding_window_possible_cached_prefix():

114 sliding_window=4,

115 )

116

114 - block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)

117 + block_pool = BlockPool(

118 + num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size

119 + )

120 manager = get_sliding_window_manager(sliding_window_spec, block_pool)

121

122 def run_one_case(block_is_cached, expect_length):

@@ -138,6 +143,7 @@ def test_sliding_window_possible_cached_prefix():

143 block_pool=block_pool,

144 kv_cache_spec=sliding_window_spec,

145 use_eagle=False,

146 + alignment_tokens=block_size,

147 )[0]

148 assert len(computed_blocks) == expect_length

149

@@ -178,7 +184,7 @@ def test_chunked_local_attention_remove_skipped_blocks():

184 attention_chunk_size=4,

185 )

186

181 - block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)

187 + block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=2)

188

189 manager = get_chunked_local_attention_manager(attention_spec, block_pool)

190

@@ -239,7 +245,7 @@ def test_sliding_window_remove_skipped_blocks():

245 sliding_window=4,

246 )

247

242 - block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)

248 + block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=2)

249

250 manager = get_sliding_window_manager(sliding_window_spec, block_pool)

251

@@ -316,7 +322,9 @@ def test_get_num_blocks_to_allocate():

322 sliding_window=4, # Placeholder value, not related to test result

323 )

324

319 - block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)

325 + block_pool = BlockPool(

326 + num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size

327 + )

328 manager = get_sliding_window_manager(sliding_window_spec, block_pool)

329 cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]

330 cached_blocks_2 = [block_pool.null_block for _ in range(5)] + [

@@ -341,7 +349,9 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():

349 attention_chunk_size=4, # Placeholder value, not related to test result

350 )

351

344 - block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)

352 + block_pool = BlockPool(

353 + num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size

354 + )

355 manager = get_chunked_local_attention_manager(attention_spec, block_pool)

356 cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]

357 cached_blocks_2 = [block_pool.null_block for _ in range(5)] + [

vllm/engine/arg_utils.py

+5 -3

@@ -1816,9 +1816,11 @@ class EngineArgs:

1816 if model_config.runner_type != "pooling":

1817 default_chunked_prefill = True

1818

1819 - # Disable prefix caching default for hybrid models

1820 - # since the feature is still experimental.

1821 - default_prefix_caching = not model_config.is_hybrid

1819 + # Disable prefix caching default for hybrid models and mamba-only

1820 + # models since the feature is still experimental.

1821 + default_prefix_caching = not (

1822 + model_config.is_hybrid or model_config.is_attention_free

1823 + )

1824 else:

1825 assert model_config.pooler_config is not None

1826

vllm/model_executor/models/config.py

+8 -3

@@ -289,9 +289,6 @@ class MambaModelConfig(VerifyAndUpdateConfig):

289 model_config = vllm_config.model_config

290 cache_config = vllm_config.cache_config

291

292 - if cache_config.mamba_block_size is None:

293 - cache_config.mamba_block_size = model_config.max_model_len

294 -

292 if cache_config.enable_prefix_caching:

293 if model_config.supports_mamba_prefix_caching:

294 logger.info(

@@ -299,6 +296,11 @@ class MambaModelConfig(VerifyAndUpdateConfig):

296 "Its support for Mamba layers is experimental. "

297 "Please report any issues you may observe."

298 )

299 + # By default, mamba block size will be set to max_model_len (see

300 + # below). When enabling prefix caching, we align mamba block size

301 + # to the block size as the basic granularity for prefix caching.

302 + if cache_config.mamba_block_size is None:

303 + cache_config.mamba_block_size = cache_config.block_size

304 else:

305 logger.info(

306 "Hybrid or mamba-based model detected without "

@@ -306,6 +308,9 @@ class MambaModelConfig(VerifyAndUpdateConfig):

308 )

309 cache_config.enable_prefix_caching = False

310

311 + if cache_config.mamba_block_size is None:

312 + cache_config.mamba_block_size = model_config.max_model_len

313 +

314 # TODO(tdoublep): remove once cascade attention is supported

315 logger.info(

316 "Disabling cascade attention since it is not supported for hybrid models."

vllm/v1/core/block_pool.py

+21 -1

@@ -13,6 +13,8 @@ from vllm.distributed.kv_events import (

13 from vllm.logger import init_logger

14 from vllm.v1.core.kv_cache_utils import (

15 BlockHash,

16 + BlockHashList,

17 + BlockHashListWithBlockSize,

18 BlockHashWithGroupId,

19 ExternalBlockHash,

20 FreeKVCacheBlockQueue,

@@ -133,6 +135,10 @@ class BlockPool:

135 Args:

136 num_gpu_blocks: The number of blocks in the pool.

137 enable_caching: Whether to enable prefix caching.

138 + hash_block_size: The block size of which the block hashes are computed.

139 + The actual block size usually equals hash_block_size, but in cases

140 + where different KV cache groups have different block sizes, the

141 + actual block size can be a multiple of hash_block_size.

142 enable_kv_cache_events: Whether to enable kv cache events.

143 """

144

@@ -140,11 +146,13 @@ class BlockPool:

146 self,

147 num_gpu_blocks: int,

148 enable_caching: bool,

149 + hash_block_size: int,

150 enable_kv_cache_events: bool = False,

151 ):

152 assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0

153 self.num_gpu_blocks = num_gpu_blocks

154 self.enable_caching = enable_caching

155 + self.hash_block_size = hash_block_size

156 # All kv-cache blocks.

157 self.blocks: list[KVCacheBlock] = [

158 KVCacheBlock(idx) for idx in range(num_gpu_blocks)

@@ -223,8 +231,20 @@ class BlockPool:

231 return

232 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]

233 assert len(request.block_hashes) >= num_full_blocks

226 - new_block_hashes = request.block_hashes[num_cached_blocks:]

234 + if block_size == self.hash_block_size:

235 + # Common case.

236 + block_hashes: BlockHashList = request.block_hashes

237 + else:

238 + # block_size is a multiple of hash_block_size. This happens when

239 + # different KV cache groups have different block sizes.

240 + assert block_size % self.hash_block_size == 0

241 + # Recalculate block_hashes at the granularity of block_size, using

242 + # the original block_hashes (at the granularity of hash_block_size).

243 + block_hashes = BlockHashListWithBlockSize(

244 + request.block_hashes, self.hash_block_size, block_size

245 + )

246

247 + new_block_hashes = block_hashes[num_cached_blocks:]

248 new_hashes: list[ExternalBlockHash] | None = (

249 [] if self.enable_kv_cache_events else None

250 )

vllm/v1/core/kv_cache_coordinator.py

+79 -19

@@ -2,15 +2,25 @@

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

3 from abc import ABC, abstractmethod

4 from collections.abc import Sequence

5 +from math import lcm

7 from vllm.v1.core.block_pool import BlockPool

7 -from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

8 +from vllm.v1.core.kv_cache_utils import (

9 + BlockHash,

10 + BlockHashList,

11 + BlockHashListWithBlockSize,

12 + KVCacheBlock,

13 +)

14 from vllm.v1.core.single_type_kv_cache_manager import (

15 CrossAttentionManager,

16 FullAttentionManager,

17 get_manager_for_kv_cache_spec,

18 )

13 -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheSpec

19 +from vllm.v1.kv_cache_interface import (

20 + FullAttentionSpec,

21 + KVCacheConfig,

22 + KVCacheSpec,

23 +)

24 from vllm.v1.request import Request

@@ -28,13 +38,17 @@ class KVCacheCoordinator(ABC):

38 enable_kv_cache_events: bool,

39 dcp_world_size: int,

40 pcp_world_size: int,

41 + hash_block_size: int,

42 ):

43 self.kv_cache_config = kv_cache_config

44 self.max_model_len = max_model_len

45 self.enable_caching = enable_caching

47 self.block_pool = BlockPool(

37 - kv_cache_config.num_blocks, enable_caching, enable_kv_cache_events

48 + kv_cache_config.num_blocks,

49 + enable_caching,

50 + hash_block_size,

51 + enable_kv_cache_events,

52 )

54 # Needs special handling for find_longest_cache_hit if eagle is enabled

@@ -213,6 +227,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

227 enable_kv_cache_events: bool,

228 dcp_world_size: int,

229 pcp_world_size: int,

230 + hash_block_size: int,

231 ):

232 super().__init__(

233 kv_cache_config,

@@ -222,6 +237,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

237 enable_kv_cache_events,

238 dcp_world_size=dcp_world_size,

239 pcp_world_size=pcp_world_size,

240 + hash_block_size=hash_block_size,

241 )

242 self.num_single_type_manager = len(self.single_type_managers)

243

@@ -255,6 +271,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

271 enable_kv_cache_events: bool,

272 dcp_world_size: int,

273 pcp_world_size: int,

274 + hash_block_size: int,

275 ):

276 super().__init__(

277 kv_cache_config,

@@ -264,6 +281,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

281 enable_kv_cache_events,

282 dcp_world_size=dcp_world_size,

283 pcp_world_size=pcp_world_size,

284 + hash_block_size=hash_block_size,

285 )

286 self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec

287 self.block_size = self.kv_cache_spec.block_size

@@ -273,6 +291,11 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

291 self.block_size *= dcp_world_size

292 if pcp_world_size > 1:

293 self.block_size *= pcp_world_size

294 + # For models using only Mamba, block_size is set to max_model_len when

295 + # prefix caching is disabled, and hash_block_size validation is skipped.

296 + assert not enable_caching or (hash_block_size == self.block_size), (

297 + "UnitaryKVCacheCoordinator assumes hash_block_size == block_size"

298 + )

299 assert len(self.kv_cache_config.kv_cache_groups) == 1, (

300 "UnitaryKVCacheCoordinator assumes only one kv cache group"

301 )

@@ -289,6 +312,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

312 block_pool=self.block_pool,

313 kv_cache_spec=self.kv_cache_spec,

314 use_eagle=self.use_eagle,

315 + alignment_tokens=self.block_size,

316 dcp_world_size=self.dcp_world_size,

317 pcp_world_size=self.pcp_world_size,

318 )

@@ -313,6 +337,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

337 enable_kv_cache_events: bool,

338 dcp_world_size: int,

339 pcp_world_size: int,

340 + hash_block_size: int,

341 ):

342 super().__init__(

343 kv_cache_config,

@@ -322,7 +347,17 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

347 enable_kv_cache_events,

348 dcp_world_size=dcp_world_size,

349 pcp_world_size=pcp_world_size,

350 + hash_block_size=hash_block_size,

351 )

352 + # hash_block_size: the block size used to compute block hashes.

353 + # The actual block size usually equals hash_block_size, but in cases where

354 + # different KV cache groups have different block sizes, the actual block size

355 + # can be a multiple of hash_block_size.

356 + self.hash_block_size = hash_block_size

357 + assert all(

358 + g.kv_cache_spec.block_size % hash_block_size == 0

359 + for g in kv_cache_config.kv_cache_groups

360 + ), "block_size must be divisible by hash_block_size"

361 assert dcp_world_size == 1, "DCP not support hybrid attn now."

362 assert pcp_world_size == 1, "PCP not support hybrid attn now."

363 self.verify_and_split_kv_cache_groups()

@@ -373,14 +408,12 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

408 self.other_spec = other_spec

409 self.full_attention_block_size = self.full_attention_spec.block_size

410 self.other_block_size = self.other_spec.block_size

376 -

377 - if self.enable_caching:

378 - # this requirement is only needed for the prefix caching logic

379 - divisible = self.other_block_size % self.full_attention_block_size

380 - assert divisible == 0, (

381 - "KVCacheCoordinator assumes the block_size of full "

382 - "attention layers is divisible by other layers now."

383 - )

411 + # The LCM of the block sizes of full attention and other attention.

412 + # The cache hit length must be a multiple of the LCM of the block sizes

413 + # to make sure the cache hit length is a multiple of the block size of

414 + # each attention type. Requiring this because we don't support partial

415 + # block cache hit yet.

416 + self.lcm_block_size = lcm(self.full_attention_block_size, self.other_block_size)

417

418 if max(self.full_attention_group_ids) < min(self.other_group_ids):

419 self.full_attn_first = True

@@ -414,25 +447,48 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

447 - The number of tokens of the longest cache hit.

448 """

449 # First, find the longest cache hit for full attention.

450 + if self.full_attention_spec.block_size == self.hash_block_size:

451 + # Common case.

452 + full_attention_block_hashes: BlockHashList = block_hashes

453 + else:

454 + # block_size is a multiple of hash_block_size. This happens when different

455 + # KV cache groups have different block sizes. In this case, we need to

456 + # recalculate block_hashes at the granularity of block_size, using the

457 + # original block_hashes (at the granularity of hash_block_size).

458 + full_attention_block_hashes = BlockHashListWithBlockSize(

459 + block_hashes, self.hash_block_size, self.full_attention_spec.block_size

460 + )

461 hit_blocks_full_attn = self.full_attention_manager_cls.find_longest_cache_hit(

418 - block_hashes=block_hashes,

462 + block_hashes=full_attention_block_hashes,

463 max_length=max_cache_hit_length,

464 kv_cache_group_ids=self.full_attention_group_ids,

465 block_pool=self.block_pool,

466 kv_cache_spec=self.full_attention_spec,

467 use_eagle=self.use_eagle,

468 + alignment_tokens=self.lcm_block_size,

469 )

470 hit_length = len(hit_blocks_full_attn[0]) * self.full_attention_block_size

471

472 # Next, find the cache hit for the other attention WITHIN

473 # the cache hit of full attention.

474 + if self.other_spec.block_size == self.hash_block_size:

475 + # Common case.

476 + other_block_hashes: BlockHashList = block_hashes

477 + else:

478 + # Similar to the full attention case, here we need to recalculate

479 + # block_hashes at the granularity of block_size, using the original

480 + # block_hashes (at the granularity of hash_block_size).

481 + other_block_hashes = BlockHashListWithBlockSize(

482 + block_hashes, self.hash_block_size, self.other_spec.block_size

483 + )

484 hit_blocks_other_attn = self.other_attention_cls.find_longest_cache_hit(

430 - block_hashes=block_hashes,

485 + block_hashes=other_block_hashes,

486 max_length=hit_length,

487 kv_cache_group_ids=self.other_group_ids,

488 block_pool=self.block_pool,

489 kv_cache_spec=self.other_spec,

490 use_eagle=self.use_eagle,

491 + alignment_tokens=self.lcm_block_size,

492 )

493 hit_length = len(hit_blocks_other_attn[0]) * self.other_block_size

494

@@ -466,6 +522,7 @@ def get_kv_cache_coordinator(

522 enable_kv_cache_events: bool,

523 dcp_world_size: int,

524 pcp_world_size: int,

525 + hash_block_size: int,

526 ) -> KVCacheCoordinator:

527 if not enable_caching:

528 return KVCacheCoordinatorNoPrefixCache(

@@ -473,8 +530,9 @@ def get_kv_cache_coordinator(

530 max_model_len,

531 use_eagle,

532 enable_kv_cache_events,

476 - dcp_world_size=dcp_world_size,

477 - pcp_world_size=pcp_world_size,

533 + dcp_world_size,

534 + pcp_world_size,

535 + hash_block_size,

536 )

537 if len(kv_cache_config.kv_cache_groups) == 1:

538 return UnitaryKVCacheCoordinator(

@@ -483,8 +541,9 @@ def get_kv_cache_coordinator(

541 use_eagle,

542 enable_caching,

543 enable_kv_cache_events,

486 - dcp_world_size=dcp_world_size,

487 - pcp_world_size=pcp_world_size,

544 + dcp_world_size,

545 + pcp_world_size,

546 + hash_block_size,

547 )

548 return HybridKVCacheCoordinator(

549 kv_cache_config,

@@ -492,6 +551,7 @@ def get_kv_cache_coordinator(

551 use_eagle,

552 enable_caching,

553 enable_kv_cache_events,

495 - dcp_world_size=dcp_world_size,

496 - pcp_world_size=pcp_world_size,

554 + dcp_world_size,

555 + pcp_world_size,

556 + hash_block_size,

557 )

vllm/v1/core/kv_cache_manager.py

+5 -20

@@ -95,6 +95,7 @@ class KVCacheManager:

95 self,

96 kv_cache_config: KVCacheConfig,

97 max_model_len: int,

98 + hash_block_size: int,

99 enable_caching: bool = True,

100 use_eagle: bool = False,

101 log_stats: bool = False,

@@ -107,28 +108,11 @@ class KVCacheManager:

108 self.enable_caching = enable_caching

109 self.use_eagle = use_eagle

110 self.log_stats = log_stats

110 - # FIXME: make prefix cache stats conditional on log_stats

111 + # FIXME: make prefix cache stats conditional on log_stats. We still need

112 + # this comment because when the log stats is enabled there are still

113 + # potential configs we could expose in the future.

114 self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

115

113 - self.block_size: int | None = None

114 - if self.enable_caching:

115 - assert (

116 - len(

117 - set(

118 - g.kv_cache_spec.block_size

119 - for g in kv_cache_config.kv_cache_groups

120 - )

121 - )

122 - == 1

123 - ), "Only one block size is supported for now"

124 - self.block_size = kv_cache_config.kv_cache_groups[

125 - 0

126 - ].kv_cache_spec.block_size

127 -

128 - if dcp_world_size * pcp_world_size > 1:

129 - assert len(kv_cache_config.kv_cache_groups) == 1

130 - self.block_size *= dcp_world_size * pcp_world_size

131 -

116 self.coordinator = get_kv_cache_coordinator(

117 kv_cache_config=kv_cache_config,

118 max_model_len=self.max_model_len,

@@ -137,6 +121,7 @@ class KVCacheManager:

121 enable_kv_cache_events=enable_kv_cache_events,

122 dcp_world_size=dcp_world_size,

123 pcp_world_size=pcp_world_size,

124 + hash_block_size=hash_block_size,

125 )

126 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)

127 self.block_pool = self.coordinator.block_pool

vllm/v1/core/kv_cache_utils.py

+136 -20

@@ -5,9 +5,9 @@

5 import copy

6 import os

7 from collections import defaultdict

8 -from collections.abc import Callable, Iterable, Sequence

9 -from dataclasses import dataclass

10 -from typing import Any, NewType, TypeAlias

8 +from collections.abc import Callable, Iterable, Iterator, Sequence

9 +from dataclasses import dataclass, replace

10 +from typing import Any, NewType, TypeAlias, overload

12 from vllm import envs

13 from vllm.config import VllmConfig

@@ -825,11 +825,11 @@ def get_num_blocks(

825 return num_blocks

826

827

828 -def get_uniform_page_size(kv_cache_spec: dict[str, KVCacheSpec]) -> int:

828 +def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:

829 """

830 Get the page size of the KV cache.

831 """

832 - page_sizes = set(layer.page_size_bytes for layer in kv_cache_spec.values())

832 + page_sizes = {layer.page_size_bytes for layer in kv_cache_specs}

833 assert len(page_sizes) == 1

834 return page_sizes.pop()

835

@@ -882,6 +882,46 @@ def is_kv_cache_page_size_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool

882 return len(page_sizes) == 1

883

884

885 +def unify_kv_cache_spec_page_size(

886 + kv_cache_spec: dict[str, KVCacheSpec],

887 +) -> dict[str, KVCacheSpec]:

888 + """

889 + Unify the page size of the given KVCacheSpec. If the page size of all layers

890 + are the same, return the original KVCacheSpec. If not same, unify the page

891 + size by increasing the block size of layers with smaller page size. Raise

892 + NotImplementedError if failed to unify the page size.

893 +

894 + Args:

895 + kv_cache_spec: The KVCacheSpec of each attention layer in the model

896 +

897 + Returns:

898 + The updated KVCacheSpec with the same page_size_bytes.

899 + """

900 + page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}

901 + if len(page_sizes) <= 1:

902 + # All layers have the same page size, no need to unify.

903 + return kv_cache_spec

904 +

905 + max_page_size = max(page_sizes)

906 + new_kv_cache_spec = {}

907 + for layer_name, layer_spec in kv_cache_spec.items():

908 + if layer_spec.page_size_bytes == max_page_size:

909 + new_kv_cache_spec[layer_name] = layer_spec

910 + else:

911 + layer_page_size = layer_spec.page_size_bytes

912 + if max_page_size % layer_page_size != 0:

913 + raise NotImplementedError(

914 + "The page size of the layer is not divisible by the "

915 + "maximum page size. Cannot unify by adjusting block_size."

916 + )

917 + ratio = max_page_size // layer_page_size

918 + new_block_size = layer_spec.block_size * ratio

919 + new_spec = replace(layer_spec, block_size=new_block_size)

920 + assert new_spec.page_size_bytes == max_page_size

921 + new_kv_cache_spec[layer_name] = new_spec

922 + return new_kv_cache_spec

923 +

924 +

925 def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:

926 # kv_cache_spec is an empty dict for attention free models

927 return not kv_cache_spec

@@ -1010,7 +1050,6 @@ def _get_kv_cache_groups_uniform_page_size(

1050 def get_kv_cache_config_from_groups(

1051 vllm_config: VllmConfig,

1052 kv_cache_groups: list[KVCacheGroupSpec],

1013 - kv_cache_specs: dict[str, KVCacheSpec],

1053 available_memory: int,

1054 ) -> KVCacheConfig:

1055 """

@@ -1020,7 +1059,6 @@ def get_kv_cache_config_from_groups(

1059 Args:

1060 vllm_config: The global VllmConfig

1061 kv_cache_groups: The KV cache groups

1023 - kv_cache_specs: The KV cache spec of each attention layer in the model

1062 available_memory: Memory available for KV cache in bytes

1063 Returns:

1064 The generated KVCacheConfig

@@ -1064,7 +1102,9 @@ def get_kv_cache_config_from_groups(

1102 # full.1, sw.2: share another Tensor with size=available_memory//2

1103 group_size = max(len(group.layer_names) for group in kv_cache_groups)

1104

1067 - page_size = get_uniform_page_size(kv_cache_specs)

1105 + page_size = get_uniform_page_size(

1106 + [group.kv_cache_spec for group in kv_cache_groups]

1107 + )

1108 assert group_size > 0, "group_size must be greater than 0"

1109 num_blocks = get_num_blocks(

1110 vllm_config, group_size, available_memory, page_size

@@ -1166,7 +1206,8 @@ def get_kv_cache_groups(

1206 # This returns an empty list to allow for the KVCacheManager to handle

1207 # attention free models.

1208 return []

1169 - elif is_kv_cache_spec_uniform(kv_cache_spec):

1209 +

1210 + if is_kv_cache_spec_uniform(kv_cache_spec):

1211 # KV cache of all layers are the same, which is true for

1212 # most models. Allocate the same amount of memory for

1213 # each layer.

@@ -1176,14 +1217,16 @@ def get_kv_cache_groups(

1217 # full attention, or all layers are sliding window attention with the

1218 # same window size). Put all layers into one group.

1219 return _get_kv_cache_groups_uniform_type(uniform_spec)

1179 - elif is_kv_cache_page_size_uniform(kv_cache_spec):

1180 - # Model contains multiple attention types, but KV cache of all layers

1181 - # have the same physical memory per block per layer. Split the layers

1182 - # into groups with the same number of layers, and thus same total page

1183 - # size.

1184 - return _get_kv_cache_groups_uniform_page_size(kv_cache_spec)

1220

1186 - raise NotImplementedError

1221 + # As KVCacheManager can only allocate memory of one size, we need to unify

1222 + # the page size of the layers. For cases cannot be unified, this function

1223 + # will raise an error.

1224 + kv_cache_spec = unify_kv_cache_spec_page_size(kv_cache_spec)

1225 + # Model contains multiple attention types, but KV cache of all layers

1226 + # have the same physical memory per block per layer. Split the layers

1227 + # into groups with the same number of layers, and thus same total page

1228 + # size.

1229 + return _get_kv_cache_groups_uniform_page_size(kv_cache_spec)

1230

1231

1232 def generate_scheduler_kv_cache_config(

@@ -1327,10 +1370,7 @@ def get_kv_cache_configs(

1370 ) == len(kv_cache_spec_one_worker), "Some layers are not assigned to any group."

1371 kv_cache_configs.append(

1372 get_kv_cache_config_from_groups(

1330 - vllm_config,

1331 - kv_cache_groups_one_worker,

1332 - kv_cache_spec_one_worker,

1333 - available_memory_one_worker,

1373 + vllm_config, kv_cache_groups_one_worker, available_memory_one_worker

1374 )

1375 )

1376

@@ -1353,3 +1393,79 @@ def get_kv_cache_configs(

1393 _report_kv_cache_config(vllm_config, kv_cache_config)

1394

1395 return kv_cache_configs

1396 +

1397 +

1398 +class BlockHashListWithBlockSize:

1399 + """

1400 + Convert block-hash granularity from `hash_block_size` to `target_block_size`.

1401 + Used when KV cache groups have different block sizes: `hash_block_size`

1402 + is the size used to compute the original `block_hashes`; `target_block_size`

1403 + is the group's actual block size.

1404 +

1405 + Currently, only scaling up by an integer factor is supported (i.e.,

1406 + `target_block_size` is a multiple of `hash_block_size`). Conversion is

1407 + performed lazily on access for efficiency, by concatenating consecutive

1408 + hashes at `hash_block_size` to form each hash at `target_block_size`.

1409 +

1410 + Example (`hash_block_size` = 16, `target_block_size` = 32):

1411 + concatenating two 16-size hashes yields one 32-size hash:

1412 +

1413 + Block hashes with block_size 16:

1414 + | Token Range | 0-15 | 16-31 | 32-47 | 48-63 |

1415 + |-------------|------|-------|-------|-------|

1416 + | Hash | A | B | C | D |

1417 +

1418 + Block hashes with block_size 32:

1419 + | Token Range | 0-31 | 32-63 |

1420 + |-------------|------|-------|

1421 + | Hash | AB | CD |

1422 +

1423 + Args:

1424 + block_hashes: Block hashes to convert, computed at `hash_block_size`.

1425 + hash_block_size: Block size at which `block_hashes` were computed.

1426 + target_block_size: Desired block size; must be a multiple of `hash_block_size`.

1427 + """

1428 +

1429 + def __init__(

1430 + self,

1431 + block_hashes: list[BlockHash],

1432 + hash_block_size: int,

1433 + target_block_size: int,

1434 + ):

1435 + self.block_hashes = block_hashes

1436 + assert target_block_size % hash_block_size == 0

1437 + self.scale_factor = target_block_size // hash_block_size

1438 +

1439 + def __len__(self) -> int:

1440 + return len(self.block_hashes) // self.scale_factor

1441 +

1442 + @overload

1443 + def __getitem__(self, idx: int) -> BlockHash: ...

1444 +

1445 + @overload

1446 + def __getitem__(self, idx: slice) -> list[BlockHash]: ...

1447 +

1448 + def __getitem__(self, idx):

1449 + if isinstance(idx, int):

1450 + return self._get_value_at(idx)

1451 +

1452 + if isinstance(idx, slice):

1453 + start, stop, step = idx.indices(len(self))

1454 + return [self._get_value_at(i) for i in range(start, stop, step)]

1455 +

1456 + raise TypeError(f"Invalid index type: {type(idx)!r}")

1457 +

1458 + def __iter__(self) -> Iterator[BlockHash]:

1459 + for i in range(len(self)):

1460 + yield self._get_value_at(i)

1461 +

1462 + def _get_value_at(self, idx: int) -> BlockHash:

1463 + base = idx * self.scale_factor

1464 + end = base + self.scale_factor

1465 + merged_hash: bytes = self.block_hashes[base]

1466 + for i in range(base + 1, end):

1467 + merged_hash += self.block_hashes[i]

1468 + return BlockHash(merged_hash)

1469 +

1470 +

1471 +BlockHashList = list[BlockHash] | BlockHashListWithBlockSize

vllm/v1/core/sched/scheduler.py

@@ -186,6 +186,7 @@ class Scheduler(SchedulerInterface):

186 enable_kv_cache_events=self.enable_kv_cache_events,

187 dcp_world_size=self.dcp_world_size,

188 pcp_world_size=self.pcp_world_size,

189 + hash_block_size=self.block_size,

190 )

191 self.use_pp = self.parallel_config.pipeline_parallel_size > 1

192 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER

vllm/v1/core/single_type_kv_cache_manager.py

+63 -9

@@ -7,7 +7,7 @@ from collections.abc import Sequence

8 from vllm.utils.math_utils import cdiv

9 from vllm.v1.core.block_pool import BlockPool

10 -from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

10 +from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock

11 from vllm.v1.kv_cache_interface import (

12 ChunkedLocalAttentionSpec,

13 CrossAttentionSpec,

@@ -207,12 +207,13 @@ class SingleTypeKVCacheManager(ABC):

207 @abstractmethod

208 def find_longest_cache_hit(

209 cls,

210 - block_hashes: list[BlockHash],

210 + block_hashes: BlockHashList,

211 max_length: int,

212 kv_cache_group_ids: list[int],

213 block_pool: BlockPool,

214 kv_cache_spec: KVCacheSpec,

215 use_eagle: bool,

216 + alignment_tokens: int,

217 dcp_world_size: int = 1,

218 pcp_world_size: int = 1,

219 ) -> tuple[list[KVCacheBlock], ...]:

@@ -232,6 +233,11 @@ class SingleTypeKVCacheManager(ABC):

233 block_pool: The block pool.

234 kv_cache_spec: The kv cache spec.

235 use_eagle: Whether to use eagle.

236 + alignment_tokens: The returned cache hit length (in tokens) should

237 + be a multiple of this value (in tokens). By default, it should

238 + be set to the block_size.

239 + dcp_world_size: The world size of decode context parallelism.

240 + pcp_world_size: The world size of prefill context parallelism.

241

242 Returns:

243 A list of cached blocks with skipped blocks replaced by null block

@@ -299,17 +305,18 @@ class FullAttentionManager(SingleTypeKVCacheManager):

305 @classmethod

306 def find_longest_cache_hit(

307 cls,

302 - block_hashes: list[BlockHash],

308 + block_hashes: BlockHashList,

309 max_length: int,

310 kv_cache_group_ids: list[int],

311 block_pool: BlockPool,

312 kv_cache_spec: KVCacheSpec,

313 use_eagle: bool,

314 + alignment_tokens: int,

315 dcp_world_size: int = 1,

316 pcp_world_size: int = 1,

317 ) -> tuple[list[KVCacheBlock], ...]:

318 assert isinstance(

312 - kv_cache_spec, (FullAttentionSpec, ChunkedLocalAttentionSpec)

319 + kv_cache_spec, FullAttentionSpec | ChunkedLocalAttentionSpec

320 ), (

321 "FullAttentionManager can only be used for full attention "

322 "and chunked local attention groups"

@@ -333,6 +340,13 @@ class FullAttentionManager(SingleTypeKVCacheManager):

340 else:

341 break

342 if use_eagle and computed_blocks[0]:

343 + # Need to drop the last matched block if eagle is enabled.

344 + for computed in computed_blocks:

345 + computed.pop()

346 + while (

347 + block_size != alignment_tokens # Faster for common case.

348 + and len(computed_blocks[0]) * block_size % alignment_tokens != 0

349 + ):

350 for computed in computed_blocks:

351 computed.pop()

352 return computed_blocks

@@ -359,12 +373,13 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

373 @classmethod

374 def find_longest_cache_hit(

375 cls,

362 - block_hashes: list[BlockHash],

376 + block_hashes: BlockHashList,

377 max_length: int,

378 kv_cache_group_ids: list[int],

379 block_pool: BlockPool,

380 kv_cache_spec: KVCacheSpec,

381 use_eagle: bool,

382 + alignment_tokens: int,

383 dcp_world_size: int = 1,

384 pcp_world_size: int = 1,

385 ) -> tuple[list[KVCacheBlock], ...]:

@@ -396,6 +411,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

411 [block_pool.null_block] * max_num_blocks

412 for _ in range(len(kv_cache_group_ids))

413 )

414 + block_size = kv_cache_spec.block_size

415 num_contiguous_blocks = 0

416 match_found = False

417 # Search from right to left and early stop when a match is found.

@@ -403,6 +419,15 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

419 if cached_block := block_pool.get_cached_block(

420 block_hashes[i], kv_cache_group_ids

421 ):

422 + # Skip prefix matching check if the block is not aligned with

423 + # `alignment_tokens`.

424 + if (

425 + num_contiguous_blocks == 0

426 + and block_size != alignment_tokens # Faster for common case.

427 + and (i + 1) * block_size % alignment_tokens != 0

428 + ):

429 + continue

430 + # Add the cached block to the computed blocks.

431 for computed, cached in zip(computed_blocks, cached_block):

432 computed[i] = cached

433 num_contiguous_blocks += 1

@@ -421,7 +446,16 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

446 # `num_contiguous_blocks < sliding_window_contiguous_blocks`.

447 for computed in computed_blocks:

448 del computed[num_contiguous_blocks:]

449 + while (

450 + block_size != alignment_tokens # Faster for common case.

451 + and len(computed_blocks[0]) * block_size % alignment_tokens != 0

452 + ):

453 + for computed in computed_blocks:

454 + computed.pop()

455 if use_eagle and computed_blocks[0]:

456 + assert kv_cache_spec.block_size == alignment_tokens, (

457 + "aligned_length is not compatible with eagle now"

458 + )

459 for computed in computed_blocks:

460 computed.pop()

461 return computed_blocks

@@ -475,12 +509,13 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

509 @classmethod

510 def find_longest_cache_hit(

511 cls,

478 - block_hashes: list[BlockHash],

512 + block_hashes: BlockHashList,

513 max_length: int,

514 kv_cache_group_ids: list[int],

515 block_pool: BlockPool,

516 kv_cache_spec: KVCacheSpec,

517 use_eagle: bool,

518 + alignment_tokens: int,

519 dcp_world_size: int = 1,

520 pcp_world_size: int = 1,

521 ) -> tuple[list[KVCacheBlock], ...]:

@@ -511,6 +546,10 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

546 block_pool: The block pool.

547 kv_cache_spec: The kv cache spec.

548 use_eagle: Whether to use eagle.

549 + dcp_world_size: The world size of decode context parallelism.

550 + pcp_world_size: The world size of prefill context parallelism.

551 + alignment_tokens: The returned cache hit length (in tokens) should

552 + be a multiple of this value (in tokens).

553

554 Returns:

555 A list of cached blocks

@@ -524,6 +563,10 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

563 )

564 assert dcp_world_size == 1, "DCP not support chunked local attn now."

565 assert pcp_world_size == 1, "PCP not support chunked local attn now."

566 + assert kv_cache_spec.block_size == alignment_tokens, (

567 + "KV cache groups with different block sizes are not compatible with "

568 + "chunked local attention now"

569 + )

570 max_num_blocks = max_length // kv_cache_spec.block_size

571 if max_length > 0:

572 local_attention_start_idx = (

@@ -612,12 +655,13 @@ class MambaManager(SingleTypeKVCacheManager):

655 @classmethod

656 def find_longest_cache_hit(

657 cls,

615 - block_hashes: list[BlockHash],

658 + block_hashes: BlockHashList,

659 max_length: int,

660 kv_cache_group_ids: list[int],

661 block_pool: BlockPool,

662 kv_cache_spec: KVCacheSpec,

663 use_eagle: bool,

664 + alignment_tokens: int,

665 dcp_world_size: int = 1,

666 pcp_world_size: int = 1,

667 ) -> tuple[list[KVCacheBlock], ...]:

@@ -630,12 +674,21 @@ class MambaManager(SingleTypeKVCacheManager):

674 [] for _ in range(len(kv_cache_group_ids))

675 )

676

633 - max_num_blocks = max_length // kv_cache_spec.block_size

677 + block_size = kv_cache_spec.block_size

678 + max_num_blocks = max_length // block_size

679 # Search from right to left and early stop when a match is found.

680 for i in range(max_num_blocks - 1, -1, -1):

681 if cached_block := block_pool.get_cached_block(

682 block_hashes[i], kv_cache_group_ids

683 ):

684 + # When enable Mamba prefix caching, `block_size` will be aligned

685 + # across full attention layers and Mamba layers to ensure the

686 + # prefix hit length aligned at block

687 + if (

688 + block_size != alignment_tokens # Faster for common case.

689 + and (i + 1) * block_size % alignment_tokens != 0

690 + ):

691 + continue

692 for computed, cached in zip(computed_blocks, cached_block):

693 # the hit length logic later assumes:

694 # hit_length = len(hit_blocks_other_attn[0])

@@ -708,12 +761,13 @@ class CrossAttentionManager(SingleTypeKVCacheManager):

761 @classmethod

762 def find_longest_cache_hit(

763 cls,

711 - block_hashes: list[BlockHash],

764 + block_hashes: BlockHashList,

765 max_length: int,

766 kv_cache_group_ids: list[int],

767 block_pool: BlockPool,

768 kv_cache_spec: KVCacheSpec,

769 use_eagle: bool,

770 + alignment_tokens: int,

771 dcp_world_size: int = 1,

772 pcp_world_size: int = 1,

773 ) -> tuple[list[KVCacheBlock], ...]: