fix(dflash): isolate replicated draft cache groups

tests/v1/core/test_kv_cache_utils.py

@@ -2232,6 +2232,61 @@ def test_group_and_unify_kv_cache_specs_mixed_page_size_groups():
     assert layer_names == {"mla.0", "mla.1", "swa.0"}
 
 
+def test_group_dcp_replicated_dflash_draft():
+    target = new_mla_spec()
+    draft = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=2048,
+        dcp_replicated=True,
+    )
+    assert target.page_size_bytes != draft.page_size_bytes
+
+    specs = {"model.layers.0": target, "draft.layers.0": draft}
+    # DeepSeek-V4's UniformType tuple planner is not needed for DFlash.
+    assert group_and_unify_kv_cache_specs(specs) is None
+
+    groups = get_kv_cache_groups(_grouping_config(), specs)
+    draft_group = next(
+        group for group in groups if isinstance(group.kv_cache_spec, SlidingWindowSpec)
+    )
+    assert all(group.kv_cache_spec.block_size == 16 for group in groups)
+    assert draft_group.kv_cache_spec.dcp_replicated is True
+
+
+def test_group_dcp_replicated_dflash_draft_with_hybrid_target():
+    target_mla = new_mla_spec()
+    target_mamba = new_mamba_spec(page_size_padded=target_mla.page_size_bytes)
+    draft = SlidingWindowSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=2048,
+        dcp_replicated=True,
+    )
+    assert target_mla.page_size_bytes == target_mamba.page_size_bytes
+    assert draft.page_size_bytes % target_mla.page_size_bytes != 0
+
+    specs = {
+        "model.mla": target_mla,
+        "model.mamba": target_mamba,
+        "draft.attn": draft,
+    }
+    groups = get_kv_cache_groups(_grouping_config(), specs)
+
+    assert {name for group in groups for name in group.layer_names} == set(specs)
+    draft_group = next(group for group in groups if "draft.attn" in group.layer_names)
+    assert draft_group.kv_cache_spec == draft
+    target_groups = [group for group in groups if "draft.attn" not in group.layer_names]
+    assert all(
+        not getattr(group.kv_cache_spec, "dcp_replicated", False)
+        for group in target_groups
+    )
+
+
 def new_indexer_mla_spec(block_size=16):
     # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller
     # page size than the main MLA attention (uint8, small head), so their pages

vllm/v1/core/kv_cache_utils.py

@@ -1600,6 +1600,37 @@ def group_and_unify_kv_cache_specs(
     return [mla_uniform_spec, *swa_uniform_specs]
 
 
+def group_dcp_replicated_draft_kv_cache_specs(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+) -> list[KVCacheGroupSpec] | None:
+    """Keep replicated speculative KV separate from a sharded target.
+
+    Replicated draft cache groups have different allocation and DCP semantics
+    from target cache groups. Group each partition independently so a hybrid
+    target retains its native grouping and page sizes instead of being
+    page-size-unified with the draft.
+    """
+    replicated = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if getattr(spec, "dcp_replicated", False)
+    }
+    if not replicated:
+        return None
+    sharded = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if not getattr(spec, "dcp_replicated", False)
+    }
+    if not sharded:
+        return None
+    return [
+        *get_kv_cache_groups(vllm_config, sharded),
+        *get_kv_cache_groups(vllm_config, replicated),
+    ]
+
+
 def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
     """Pick a chunk size that minimizes total upward padding.
 
@@ -1784,6 +1815,10 @@ def get_kv_cache_groups(
         # full attention, or all layers are sliding window attention with the
         # same window size). Put all layers into one group.
         return _get_kv_cache_groups_uniform_type(uniform_spec)
+    elif replicated_groups := group_dcp_replicated_draft_kv_cache_specs(
+        vllm_config, kv_cache_spec
+    ):
+        return replicated_groups
     elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec):
         # DeepseekV4 case: All layers need the same number of token slots,
         # yet some layers are full attention while others are sliding window