Fix hybrid DCP prefix cache geometry

Author: s0l

tests/v1/core/test_dcp_mamba_chunk_alignment.py

@@ -5,6 +5,14 @@
 
 import pytest
 
+from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
+
 
 def test_mamba_align_chunk_split_uses_dcp_effective_block_size():
     pytest.importorskip("torch")
@@ -29,7 +37,8 @@ def test_mamba_align_chunk_split_uses_dcp_effective_block_size():
     )
 
     assert adjusted == 96
-    assert adjusted % (scheduler.cache_config.block_size * scheduler.dcp_world_size) == 0
+    effective_block_size = scheduler.cache_config.block_size * scheduler.dcp_world_size
+    assert adjusted % effective_block_size == 0
 
 
 def test_mamba_align_split_keeps_small_chunks_when_dcp_alignment_exceeds_budget():
@@ -55,3 +64,51 @@ def test_mamba_align_split_keeps_small_chunks_when_dcp_alignment_exceeds_budget(
     )
 
     assert adjusted == 256
+
+
+def test_hybrid_dcp_resolves_global_scheduler_block_and_fine_hash_block():
+    torch = pytest.importorskip("torch")
+    block_size = 16
+    kv_cache_config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=((1, 1),),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            block_size=block_size,
+            enable_prefix_caching=True,
+            hash_block_size=None,
+        ),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=3,
+            prefill_context_parallel_size=1,
+        ),
+        kv_transfer_config=None,
+    )
+
+    scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
+        kv_cache_config, vllm_config
+    )
+
+    assert scheduler_block_size == block_size * 3
+    assert hash_block_size == block_size

tests/v1/core/test_prefix_caching.py

@@ -1024,6 +1024,63 @@ def test_prefill_hybrid_model_mamba_align():
     manager.free(req0)
 
 
+def test_prefill_hybrid_model_mamba_align_dcp_replay():
+    """DCP hybrid cache insertion and lookup must use the same block geometry."""
+    block_size = 16
+    dcp_world_size = 3
+    effective_block_size = block_size * dcp_world_size
+    kv_cache_config = _make_hybrid_kv_cache_config(
+        block_size, 40, ["full", "mamba_align"]
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+        scheduler_block_size=effective_block_size,
+        dcp_world_size=dcp_world_size,
+    )
+
+    common_token_ids = list(range(3 * effective_block_size))
+    req0 = make_request("dcp-fill", common_token_ids, block_size, sha256)
+    computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
+    assert num_computed_tokens == 0
+    blocks = manager.allocate_slots(
+        req0, req0.num_tokens, num_computed_tokens, computed_blocks
+    )
+    assert blocks is not None
+    manager.new_step_starts()
+
+    coarse_hashes = kv_cache_utils.BlockHashListWithBlockSize(
+        req0.block_hashes, block_size, effective_block_size
+    )
+    cached_by_group = [
+        [
+            manager.block_pool.get_cached_block(block_hash, [group_id]) is not None
+            for block_hash in coarse_hashes
+        ]
+        for group_id in range(2)
+    ]
+    assert cached_by_group == [[True] * 3, [False, False, True]]
+
+    req1 = make_request("dcp-replay", common_token_ids + [101] * 5, block_size, sha256)
+    computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
+
+    assert num_computed_tokens == 3 * effective_block_size
+    assert all(len(group) == 3 for group in computed_blocks.blocks)
+
+    per_group_blocks, per_group_hit_lengths = (
+        manager.coordinator.find_longest_cache_hit_per_group(
+            req1.block_hashes, req1.num_tokens - 1
+        )
+    )
+    assert per_group_hit_lengths == (3 * effective_block_size,) * 2
+    assert all(len(group) == 3 for group in per_group_blocks)
+
+    manager.free(req0)
+    manager.free(req1)
+
+
 def test_hybrid_cache_mamba_align_shared_prefix_detection():
     """Test shared prefix detection heuristic for mamba align cache mode
 

vllm/v1/core/kv_cache_coordinator.py

@@ -563,6 +563,8 @@ def __init__(
                 f"share the same block size, got {group_block_sizes}."
             )
         assert pcp_world_size == 1, "PCP not support hybrid attn now."
+        self.dcp_world_size = dcp_world_size
+        self.pcp_world_size = pcp_world_size
         self.verify_and_split_kv_cache_groups()
 
     def verify_and_split_kv_cache_groups(self) -> None:
@@ -658,11 +660,15 @@ def find_longest_cache_hit(
                 - The number of tokens of the longest cache hit.
         """
 
+        def _effective_block_size(kv_cache_spec: KVCacheSpec) -> int:
+            return kv_cache_spec.block_size * self.dcp_world_size * self.pcp_world_size
+
         def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
-            if kv_cache_spec.block_size == self.hash_block_size:
+            effective_block_size = _effective_block_size(kv_cache_spec)
+            if effective_block_size == self.hash_block_size:
                 return block_hashes
             return BlockHashListWithBlockSize(
-                block_hashes, self.hash_block_size, kv_cache_spec.block_size
+                block_hashes, self.hash_block_size, effective_block_size
             )
 
         num_groups = len(self.kv_cache_config.kv_cache_groups)
@@ -687,13 +693,14 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
             for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(
                 self.attention_groups
             ):
+                effective_block_size = _effective_block_size(spec)
                 cached_blocks = hit_blocks_by_group[group_ids[0]]
                 if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:
                     # Full attention is downward-closed: we only need to look
                     # up cached blocks once; on subsequent iterations just trim
                     # to the (reduced) current hit length.
                     curr_hit_length = (
-                        curr_hit_length // spec.block_size * spec.block_size
+                        curr_hit_length // effective_block_size * effective_block_size
                     )
                     continue
 
@@ -703,7 +710,8 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
                 if drop_eagle_block:
                     # Eagle needs to match one more block and then pop the last.
                     _max_length = min(
-                        curr_hit_length + spec.block_size, max_cache_hit_length
+                        curr_hit_length + effective_block_size,
+                        max_cache_hit_length,
                     )
                 hit_blocks = manager_cls.find_longest_cache_hit(
                     block_hashes=_get_block_hashes(spec),
@@ -713,8 +721,10 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
                     kv_cache_spec=spec,
                     drop_eagle_block=drop_eagle_block,
                     alignment_tokens=self.scheduler_block_size,
+                    dcp_world_size=self.dcp_world_size,
+                    pcp_world_size=self.pcp_world_size,
                 )
-                _new_hit_length = len(hit_blocks[0]) * spec.block_size
+                _new_hit_length = len(hit_blocks[0]) * effective_block_size
                 if drop_eagle_block:
                     eagle_verified.add(idx)
                 elif _new_hit_length < curr_hit_length:
@@ -735,7 +745,7 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
         # Truncate full attention blocks to final hit_length (if present)
         first_group = self.attention_groups[0]
         if isinstance(first_group.spec, FullAttentionSpec):
-            num_blocks = hit_length // first_group.spec.block_size
+            num_blocks = hit_length // _effective_block_size(first_group.spec)
             for group_id in first_group.group_ids:
                 if (blks := hit_blocks_by_group[group_id]) is not None:
                     del blks[num_blocks:]
@@ -758,11 +768,15 @@ def find_longest_cache_hit_per_group(
             (blocks_per_group, hit_lengths_per_group)
         """
 
+        def _effective_block_size(kv_cache_spec: KVCacheSpec) -> int:
+            return kv_cache_spec.block_size * self.dcp_world_size * self.pcp_world_size
+
         def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
-            if kv_cache_spec.block_size == self.hash_block_size:
+            effective_block_size = _effective_block_size(kv_cache_spec)
+            if effective_block_size == self.hash_block_size:
                 return block_hashes
             return BlockHashListWithBlockSize(
-                block_hashes, self.hash_block_size, kv_cache_spec.block_size
+                block_hashes, self.hash_block_size, effective_block_size
             )
 
         num_groups = len(self.kv_cache_config.kv_cache_groups)
@@ -778,8 +792,10 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
                 kv_cache_spec=spec,
                 drop_eagle_block=use_eagle,
                 alignment_tokens=self.scheduler_block_size,
+                dcp_world_size=self.dcp_world_size,
+                pcp_world_size=self.pcp_world_size,
             )
-            group_hit = len(blocks[0]) * spec.block_size
+            group_hit = len(blocks[0]) * _effective_block_size(spec)
             for gid, blks in zip(group_ids, blocks):
                 hit_blocks[gid] = blks
                 hit_lengths[gid] = group_hit

vllm/v1/core/kv_cache_utils.py

@@ -627,7 +627,7 @@ def resolve_kv_cache_block_sizes(
     - ``scheduler_block_size`` is the token-alignment invariant used by the
       scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:
       ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every
-      group's block size — context parallelism is not supported here.
+      group's block size, scaled by the context-parallel world size.
     - ``hash_block_size`` is the granularity at which ``Request.block_hashes``
       is computed. Single group: equals scheduler block size. Multiple groups:
       ``cache_config.hash_block_size`` override if set, else the GCD of group
@@ -652,7 +652,7 @@ def resolve_kv_cache_block_sizes(
             "support context parallelism (dcp_world_size/pcp_world_size > 1). "
             f"Got group block sizes={group_block_sizes}."
         )
-    scheduler_block_size = math.lcm(*group_block_sizes)
+    scheduler_block_size = math.lcm(*group_block_sizes) * dcp * pcp
 
     # Block hashes are only consumed by prefix caching and KV connectors
     # (P/D, offloading); when neither is active, keep hash_block_size equal

vllm/v1/core/single_type_kv_cache_manager.py

@@ -1057,13 +1057,11 @@ def find_longest_cache_hit(
         assert isinstance(kv_cache_spec, MambaSpec), (
             "MambaManager can only be used for mamba groups"
         )
-        assert dcp_world_size == 1, "DCP not support mamba now."
-        assert pcp_world_size == 1, "PCP not support mamba now."
         computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
             [] for _ in range(len(kv_cache_group_ids))
         )
 
-        block_size = kv_cache_spec.block_size
+        block_size = kv_cache_spec.block_size * dcp_world_size * pcp_world_size
         max_num_blocks = max_length // block_size
         # Search from right to left and early stop when a match is found.
         for i in range(max_num_blocks - 1, -1, -1):