fix(dflash): restore aligned draft cache geometry

tests/v1/core/test_kv_cache_utils.py

@@ -55,6 +55,7 @@
     KVQuantMode,
     MambaSpec,
     MLAAttentionSpec,
+    RSWASpec,
     SinkFullAttentionSpec,
     SlidingWindowMLASpec,
     SlidingWindowSpec,
@@ -2079,6 +2080,30 @@ def test_get_kv_cache_configs_attention_free():
     ]
 
 
+def test_get_kv_cache_configs_preserves_model_sliding_window_retention():
+    """Generic spec-decode planning must not erase model-specific retention."""
+    model_config = ModelConfig(max_model_len=4096)
+    vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
+    vllm_config.cache_config.prefix_cache_retention_interval = None
+    spec = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=2048,
+        extra_retained_tokens=2048,
+    )
+
+    configs = get_kv_cache_configs(
+        vllm_config,
+        [{"draft": spec}],
+        [spec.page_size_bytes * 1024],
+    )
+
+    assert configs[0].kv_cache_groups[0].kv_cache_spec.extra_retained_tokens == 2048
+
+
 def test_generate_uniform_type_kv_cache_specs():
     # All layers are full attention, can be merged
     kv_cache_specs = {
@@ -2232,6 +2257,62 @@ 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_with_hybrid_mla_target():
+    target_full = new_mla_spec(block_size=16)
+    target_swa = SlidingWindowMLASpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.bfloat16,
+        sliding_window=2048,
+    )
+    draft = SlidingWindowSpec(
+        block_size=256,
+        num_kv_heads=4,
+        head_size=128,
+        dtype=torch.bfloat16,
+        sliding_window=2048,
+        dcp_replicated=True,
+    )
+    config = _grouping_config()
+    groups = get_kv_cache_groups(
+        config,
+        {"target.full": target_full, "target.swa": target_swa, "draft": draft},
+    )
+
+    assert len(groups) == 3
+    assert [group.layer_names for group in groups] == [
+        ["target.full"],
+        ["target.swa"],
+        ["draft"],
+    ]
+    assert groups[-1].kv_cache_spec.dcp_replicated is True
+
+
 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
@@ -2251,6 +2332,191 @@ def _grouping_config():
     )
 
 
+@pytest.mark.parametrize(
+    "spec",
+    [
+        MLAAttentionSpec(
+            block_size=16,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=torch.float16,
+            dcp_replicated=True,
+        ),
+        RSWASpec(
+            block_size=16,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=torch.float16,
+            dcp_replicated=True,
+            rswa_window=64,
+        ),
+        SinkFullAttentionSpec(
+            block_size=16,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=torch.float16,
+            dcp_replicated=True,
+            sink_len=4,
+        ),
+        SlidingWindowMLASpec(
+            block_size=16,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=torch.float16,
+            dcp_replicated=True,
+            sliding_window=64,
+        ),
+    ],
+    ids=("mla", "rswa", "sink", "sliding-window-mla"),
+)
+def test_attention_spec_merge_preserves_dcp_replicated(spec):
+    merged = type(spec).merge([spec, spec])
+
+    assert merged.dcp_replicated is True
+
+
+def test_sliding_window_mla_uniformity_includes_dcp_replication():
+    replicated = SlidingWindowMLASpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        dcp_replicated=True,
+        sliding_window=64,
+    )
+    sharded = SlidingWindowMLASpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        dcp_replicated=False,
+        sliding_window=64,
+    )
+
+    assert not replicated.is_uniform_with_collection(
+        {"replicated": replicated, "sharded": sharded}
+    )
+
+
+def test_disable_hybrid_manager_skips_dflash_partition():
+    target = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+    )
+    draft = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=64,
+    )
+    config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=True),
+        speculative_config=SimpleNamespace(method="dflash"),
+        model_config=SimpleNamespace(get_num_layers=lambda _parallel_config: 1),
+        parallel_config=SimpleNamespace(pipeline_parallel_size=1),
+    )
+    specs = {
+        "model.layers.0.self_attn": target,
+        "model.layers.1.self_attn": draft,
+    }
+
+    groups = get_kv_cache_groups(config, specs)
+
+    assert len(groups) == 1
+    assert set(groups[0].layer_names) == set(specs)
+
+
+def test_disable_hybrid_manager_rejects_mixed_dcp_replication():
+    target = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+    )
+    draft = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=64,
+        dcp_replicated=True,
+    )
+    config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=True),
+        speculative_config=SimpleNamespace(method="dflash"),
+        model_config=SimpleNamespace(get_num_layers=lambda _parallel_config: 1),
+        parallel_config=SimpleNamespace(pipeline_parallel_size=1),
+    )
+
+    with pytest.raises(ValueError, match="DCP-replicated and DCP-sharded"):
+        get_kv_cache_groups(
+            config,
+            {
+                "model.layers.0.self_attn": target,
+                "model.layers.1.self_attn": draft,
+            },
+        )
+
+
+def test_dflash_draft_cache_partition_is_pp1_only():
+    mamba = MambaSpec(
+        block_size=16,
+        shapes=((18432,),),
+        dtypes=(torch.bfloat16,),
+        mamba_cache_mode="align",
+    )
+    mla = MLAAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.float32,
+    )
+    draft = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.bfloat16,
+        sliding_window=2048,
+    )
+    assert len({mamba.page_size_bytes, mla.page_size_bytes, draft.page_size_bytes}) == 1
+
+    specs = {}
+    for layer_index in [i for i in range(45) if i % 4 != 3]:
+        specs[f"model.layers.{layer_index}.linear_attn"] = mamba
+    for layer_index in [i for i in range(45) if i % 4 == 3]:
+        specs[f"model.layers.{layer_index}.self_attn"] = mla
+    draft_names = {
+        f"model.layers.{layer_index}.self_attn" for layer_index in range(45, 50)
+    }
+    for layer_name in draft_names:
+        specs[layer_name] = draft
+
+    cases = ((2, 23, 11), (1, 45, 6))
+    for pipeline_parallel_size, target_layers, expected_groups in cases:
+        config = SimpleNamespace(
+            scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
+            speculative_config=SimpleNamespace(
+                method="dflash",
+                attention_backend="FLASH_ATTN",
+            ),
+            model_config=SimpleNamespace(
+                get_num_layers=lambda parallel_config,
+                target_layers=target_layers: target_layers
+            ),
+            parallel_config=SimpleNamespace(
+                pipeline_parallel_size=pipeline_parallel_size
+            ),
+        )
+        groups = get_kv_cache_groups(config, dict(specs))
+
+        assert len(groups) == expected_groups
+        if pipeline_parallel_size == 1:
+            assert any(set(group.layer_names) == draft_names for group in groups)
+
+
 def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity():
     block_size = 544
     full_spec = FullAttentionSpec(