[DCP][Mooncake] Harden hybrid prefix-cache external hits

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -78,31 +78,47 @@ def make_full_mamba_manager(
     num_blocks: int = 32,
     use_eagle: bool = False,
     num_prefill_checkpoint_blocks: int = 0,
+    include_draft_eagle_group: bool = False,
 ):
-    kv_cache_config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
+    kv_cache_groups = [
+        KVCacheGroupSpec(
+            ["full"],
+            FullAttentionSpec(
+                block_size=full_block_size,
+                num_kv_heads=1,
+                head_size=1,
+                dtype=torch.float32,
+            ),
+        ),
+        KVCacheGroupSpec(
+            ["mamba"],
+            MambaSpec(
+                block_size=mamba_block_size,
+                shapes=(1, 1),
+                dtypes=(torch.float32,),
+                mamba_cache_mode="align",
+                num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
+            ),
+        ),
+    ]
+    if include_draft_eagle_group:
+        kv_cache_groups.append(
             KVCacheGroupSpec(
-                ["full"],
+                ["draft"],
                 FullAttentionSpec(
                     block_size=full_block_size,
                     num_kv_heads=1,
                     head_size=1,
                     dtype=torch.float32,
+                    non_causal=True,
                 ),
-            ),
-            KVCacheGroupSpec(
-                ["mamba"],
-                MambaSpec(
-                    block_size=mamba_block_size,
-                    shapes=(1, 1),
-                    dtypes=(torch.float32,),
-                    mamba_cache_mode="align",
-                    num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
-                ),
-            ),
-        ],
+                is_eagle_group=True,
+            )
+        )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[],
+        kv_cache_groups=kv_cache_groups,
     )
     scheduler_block_size = lcm(
         full_block_size * dcp_world_size,
@@ -119,6 +135,122 @@ def make_full_mamba_manager(
     )
 
 
+def test_dcp_full_attention_enables_partial_hash_hits():
+    hash_block_size = 2
+    manager = make_full_mamba_manager(
+        dcp_world_size=2,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=hash_block_size,
+        num_blocks=64,
+    )
+
+    assert manager.coordinator.enable_partial_hash_hits
+    assert manager.coordinator._cache_hit_alignment_tokens == hash_block_size
+
+
+def test_dcp_fine_hit_retention_uses_hash_alignment_for_mamba():
+    """DCP must not floor Mamba's replay checkpoint to the enlarged full-attn
+    scheduler block.
+
+    Full-attention's effective block is 8 tokens under DCP4 while replicated
+    Mamba and hashing stay at 2. With latest-only retention, a 7-token prompt
+    has a reusable 6-token boundary. Flooring that boundary to 8 drops the
+    only Mamba checkpoint and collapses the reconciled hybrid hit to zero.
+    """
+    hash_block_size = 2
+    manager = make_full_mamba_manager(
+        dcp_world_size=4,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=hash_block_size,
+        num_blocks=64,
+    )
+    manager.coordinator.retention_interval = 0
+
+    token_ids = list(range(7))
+    req0 = make_request("0", token_ids, hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    assert manager.allocate_slots(req0, 6, 0, computed_blocks) is not None
+    manager.free(req0)
+    manager.new_step_starts()
+
+    req1 = make_request("1", token_ids, hash_block_size, sha256)
+    _, num_computed, _ = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+
+
+@pytest.mark.parametrize(
+    ("hash_block_size", "expected_hit", "expected_chunk_ends"),
+    [
+        (128, 6144, [4608, 6144, 7552, 7621]),
+        (1536, 4608, [4608, 7621]),
+    ],
+)
+def test_dcp_eagle_retention_primes_first_mamba_consumer(
+    hash_block_size: int,
+    expected_hit: int,
+    expected_chunk_ends: list[int],
+):
+    """The first consumer reuses the EAGLE-adjusted Mamba checkpoint."""
+    block_size = 1536
+    manager = make_full_mamba_manager(
+        dcp_world_size=8,
+        hash_block_size=hash_block_size,
+        full_block_size=block_size,
+        mamba_block_size=block_size,
+        num_blocks=128,
+        use_eagle=True,
+        include_draft_eagle_group=True,
+    )
+    manager.coordinator.retention_interval = 0
+    scheduler = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=block_size),
+        max_num_scheduled_tokens=8192,
+        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
+        use_eagle=True,
+        hash_block_size=hash_block_size,
+        mamba_partial_cache_hit=hash_block_size < block_size,
+        mamba_has_prefill_checkpoint_blocks=False,
+    )
+
+    token_ids = list(range(7621))
+    producer = make_request("producer", token_ids, hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(producer)
+    assert num_computed == 0
+    chunk_ends = []
+    while producer.num_computed_tokens < len(token_ids):
+        num_new_tokens = Scheduler._mamba_block_aligned_split(
+            scheduler,
+            producer,
+            len(token_ids) - producer.num_computed_tokens,
+        )
+        assert num_new_tokens > 0
+        assert (
+            manager.allocate_slots(
+                producer,
+                num_new_tokens,
+                num_computed,
+                computed_blocks,
+            )
+            is not None
+        )
+        producer.num_computed_tokens += num_new_tokens
+        chunk_ends.append(producer.num_computed_tokens)
+        computed_blocks = None
+        num_computed = 0
+        manager.new_step_starts()
+    assert chunk_ends == expected_chunk_ends
+
+    manager.free(producer)
+    manager.new_step_starts()
+
+    first_consumer = make_request("consumer", token_ids, hash_block_size, sha256)
+    _, num_computed, _ = manager.get_computed_blocks(first_consumer)
+    assert num_computed == expected_hit
+
+
 @pytest.mark.parametrize("dcp_world_size", [1, 4])
 def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
     """Chunk ends with partial hits on: block-aligned chunks, one extra stop
@@ -618,50 +750,42 @@ def test_external_mamba_hit_same_block_uses_running_cow_on_continue():
     assert moved[0].block_id == cow_copy.dst_block_id
 
 
-def test_take_partial_tail_offloads_returns_cow_target():
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
+def test_take_partial_tail_offloads_returns_cow_target(dcp_world_size: int):
     """The connector offload hand-off exposes the mamba CoW *target* block Y
     (the durable boundary state), not the overwritten source X, and only at
     the CoW step."""
     hash_block_size = 2
     block_size = 2 * hash_block_size
-    kv_cache_config = KVCacheConfig(
-        num_blocks=24,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["full"],
-                FullAttentionSpec(
-                    block_size=hash_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",
-                    num_prefill_checkpoint_blocks=1,
-                ),
-            ),
-        ],
-    )
-    manager = make_kv_cache_manager(
-        kv_cache_config=kv_cache_config,
-        max_model_len=8192,
-        enable_caching=True,
+    manager = make_full_mamba_manager(
+        dcp_world_size=dcp_world_size,
         hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=block_size,
+        num_blocks=24,
+        num_prefill_checkpoint_blocks=1,
     )
+    full_manager, mamba_manager = manager.coordinator.single_type_managers
+    assert full_manager.dcp_world_size == dcp_world_size
+    assert full_manager.block_size == hash_block_size * dcp_world_size
+    assert mamba_manager.dcp_world_size == 1
+    assert mamba_manager.block_size == block_size
 
     req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
     computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
     assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
 
-    # Step A registered the partial tail but has not CoW'd yet: no offload.
-    assert manager.take_partial_tail_offloads() == {}
+    # Step A registered the partial tail but has not CoW'd yet, so the mamba
+    # group has nothing durable to hand off. Under DCP the full-attention group
+    # does: its boundary block needs no CoW because KV below the boundary is
+    # append-only.
+    step_a_offloads = manager.take_partial_tail_offloads().get("0", [])
+    assert [entry for entry in step_a_offloads if entry[0] == 1] == []
+    if dcp_world_size > 1:
+        assert [entry[0] for entry in step_a_offloads] == [0]
+        assert [entry[2] for entry in step_a_offloads] == [6]
+    else:
+        assert step_a_offloads == []
 
     partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
     source_block = manager.block_pool.get_cached_block(

tests/v1/core/test_kv_cache_utils.py

@@ -1294,6 +1294,21 @@ def test_project_kv_cache_groups_to_worker():
     assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
 
 
+def test_dcp_world_size_for_kv_cache_spec_shards_full_attention_only():
+    dcp = 8
+    full = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
+    )
+    mla = new_mla_spec()
+    mamba = new_mamba_spec()
+    uniform_mla = UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs={"layer": mla})
+    assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(full, dcp) == dcp
+    assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(mla, dcp) == dcp
+    assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(uniform_mla, dcp) == dcp
+    assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(mamba, dcp) == 1
+    assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(full, 1) == 1
+
+
 @pytest.mark.parametrize(
     "layer_type,dcp_size,expected_width",
     [

tests/v1/core/test_prefix_caching.py

@@ -232,6 +232,83 @@ def make_kv_cache_config_three_types(
     )
 
 
+def test_prefix_cache_hit_uses_per_group_dcp_geometry():
+    """Prefix lookup must use each group's DCP size, not the process-wide one.
+
+    Target and draft MLA are both sharded (DCP=8); Mamba stays replicated
+    (DCP=1). Hits then align to the sharded full-attention block, not the
+    unsharded page size.
+    """
+    block_size = 16
+    dcp = 8
+    sharded_block = block_size * dcp
+    kv_cache_config = KVCacheConfig(
+        num_blocks=64,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["target_mla"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["draft_mla"],
+                MLAAttentionSpec(
+                    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,),
+                ),
+            ),
+        ],
+    )
+    manager = KVCacheManager(
+        kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+        scheduler_block_size=sharded_block,
+        dcp_world_size=dcp,
+    )
+    target_mgr, draft_mgr, mamba_mgr = manager.coordinator.single_type_managers
+    assert target_mgr.dcp_world_size == dcp
+    assert draft_mgr.dcp_world_size == dcp
+    assert mamba_mgr.dcp_world_size == 1
+    assert target_mgr.block_size == sharded_block
+    assert draft_mgr.block_size == sharded_block
+    assert mamba_mgr.block_size == block_size
+
+    hash_fn = sha256
+    common_token_ids = [i for i in range(2) for _ in range(sharded_block)]
+    req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn)
+    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
+    assert num_computed_tokens == 0
+    blocks = manager.allocate_slots(
+        req0, len(req0.prompt_token_ids), 0, computed_blocks
+    )
+    assert blocks is not None
+
+    req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn)
+    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
+    assert num_computed_tokens == 2 * sharded_block
+    assert [len(group) for group in computed_blocks.blocks] == [2, 2, 16]
+
+    manager.free(req0)
+    manager.free(req1)
+
+
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
 def test_prefill(hash_fn):
     block_size = 16