[Bugfix] Disable fine-grained prefix-cache hits for incompatible hybrid KV layouts (#51843)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 12, 2026 at 07:08 UTC
8e958902eee56ca5158728f1dd5a32246f0246f3
3 files changed
+96
-3
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+72
@@ -24,6 +24,7 @@ from vllm.v1.kv_cache_interface import (
24
KVCacheConfig,
25
KVCacheGroupSpec,
26
MambaSpec,
27
+ SlidingWindowSpec,
28
)
29
30
@@ -1353,3 +1354,74 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks():
1354
len(group) * block_size >= num_computed for group in computed_blocks.blocks
1355
)
1356
assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None
1357
+
1358
+
1359
+def test_hybrid_sliding_window_group_disables_partial_hash_hits():
1360
+ hash_block_size = 2
1361
+ sliding_window_block_size = 2 * hash_block_size
1362
+ mamba_block_size = 2 * sliding_window_block_size
1363
+ kv_cache_config = KVCacheConfig(
1364
+ num_blocks=64,
1365
+ kv_cache_tensors=[],
1366
+ kv_cache_groups=[
1367
+ KVCacheGroupSpec(
1368
+ ["full"],
1369
+ FullAttentionSpec(
1370
+ block_size=hash_block_size,
1371
+ num_kv_heads=1,
1372
+ head_size=1,
1373
+ dtype=torch.float32,
1374
+ ),
1375
+ ),
1376
+ KVCacheGroupSpec(
1377
+ ["mamba"],
1378
+ MambaSpec(
1379
+ block_size=mamba_block_size,
1380
+ shapes=(1, 1),
1381
+ dtypes=(torch.float32,),
1382
+ mamba_cache_mode="align",
1383
+ ),
1384
+ ),
1385
+ KVCacheGroupSpec(
1386
+ ["swa_draft"],
1387
+ SlidingWindowSpec(
1388
+ block_size=sliding_window_block_size,
1389
+ num_kv_heads=1,
1390
+ head_size=1,
1391
+ dtype=torch.float32,
1392
+ sliding_window=sliding_window_block_size,
1393
+ ),
1394
+ is_eagle_group=True,
1395
+ ),
1396
+ ],
1397
+ )
1398
+ manager = make_kv_cache_manager(
1399
+ kv_cache_config=kv_cache_config,
1400
+ max_model_len=8192,
1401
+ enable_caching=True,
1402
+ hash_block_size=hash_block_size,
1403
+ use_eagle=True,
1404
+ )
1405
+
1406
+ tokens = list(range(3 * sliding_window_block_size))
1407
+ request = make_request("0", tokens, hash_block_size, sha256)
1408
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
1409
+ assert not manager.coordinator.enable_partial_hash_hits
1410
+ assert (
1411
+ manager.allocate_slots(request, mamba_block_size, num_computed, computed_blocks)
1412
+ is not None
1413
+ )
1414
+ request.num_computed_tokens = mamba_block_size
1415
+ manager.new_step_starts()
1416
+ assert manager.allocate_slots(request, len(tokens) - mamba_block_size) is not None
1417
+ request.num_computed_tokens = len(tokens)
1418
+ manager.free(request)
1419
+ manager.new_step_starts()
1420
+
1421
+ cached_request = make_request(
1422
+ "1", tokens + [len(tokens), len(tokens) + 1], hash_block_size, sha256
1423
+ )
1424
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(cached_request)
1425
+
1426
+ assert num_computed == mamba_block_size
1427
+ assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed
vllm/v1/core/kv_cache_coordinator.py
+23
-3
@@ -5,6 +5,7 @@ from collections.abc import Sequence
5
from typing import NamedTuple
6
7
from vllm import envs
8
+from vllm.logger import init_logger
9
from vllm.utils.math_utils import cdiv
10
from vllm.v1.core.block_pool import BlockPool
11
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
@@ -26,6 +27,8 @@ from vllm.v1.kv_cache_interface import (
27
)
28
from vllm.v1.request import Request
29
30
+logger = init_logger(__name__)
31
+
32
33
def _validate_prefix_cache_retention_interval(
34
retention_interval: int | None,
@@ -62,6 +65,8 @@ class KVCacheCoordinator(ABC):
65
Coordinate the KV cache of different KV cache groups.
66
"""
67
68
+ enable_partial_hash_hits = False
69
+
70
def __init__(
71
self,
72
kv_cache_config: KVCacheConfig,
@@ -578,14 +583,29 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
583
"full-attention and Mamba groups, got: "
584
f"{type(g.kv_cache_spec).__name__}."
585
)
581
- # Partial hash hits are limited to full-attention + mamba ("align")
582
- # without context parallelism.
583
- self.enable_partial_hash_hits = dcp_world_size == 1 and any(
586
+ # Fine-grained hash hits require Mamba "align", no context
587
+ # parallelism, and compatible cache managers in every group.
588
+ has_partial_mamba_group = any(
589
isinstance(g.kv_cache_spec, MambaSpec)
590
and g.kv_cache_spec.mamba_cache_mode == "align"
591
and g.kv_cache_spec.block_size > hash_block_size
592
for g in kv_cache_config.kv_cache_groups
593
)
594
+ self.enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group
595
+ if self.enable_partial_hash_hits:
596
+ unsupported_partial_hit_managers = {
597
+ type(manager).__name__
598
+ for manager in self.single_type_managers
599
+ if not manager.supports_fine_grained_hash_lookup
600
+ and manager.block_size != hash_block_size
601
+ }
602
+ if unsupported_partial_hit_managers:
603
+ self.enable_partial_hash_hits = False
604
+ logger.warning_once(
605
+ "Disabling fine-grained prefix-cache hits because these KV "
606
+ "cache managers require block-aligned lookups: %s.",
607
+ ", ".join(sorted(unsupported_partial_hit_managers)),
608
+ )
609
self.verify_and_split_kv_cache_groups()
610
611
@property
vllm/v1/core/sched/scheduler.py
+1
@@ -307,6 +307,7 @@ class Scheduler(SchedulerInterface):
307
self.mamba_partial_cache_hit = (
308
self.need_mamba_block_aligned_split
309
and self.hash_block_size < self.block_size
310
+ and self.kv_cache_manager.coordinator.enable_partial_hash_hits
311
)
312
313
# Counts of non-empty steps scheduled / processed. update_from_output