fix(dflash): replicate draft KV under DCP

tests/v1/attention/test_flashinfer_dcp_spec_reorder.py

@@ -2,6 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """FlashInfer GQA builder: reorder threshold under DCP with spec decode."""
 
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 
 from vllm.platforms import current_platform
@@ -15,8 +18,10 @@
 from vllm.config import SpeculativeConfig, set_current_vllm_config
 from vllm.v1.attention.backends import flashinfer as flashinfer_backend
 from vllm.v1.attention.backends.flashinfer import (
+    BatchDCPPrefillWrapper,
     FlashInferDecodeKernel,
     FlashInferMetadataBuilder,
+    _get_dcp_local_kv_page_metadata,
 )
 from vllm.v1.attention.backends.utils import PerLayerParameters
 from vllm.v1.kv_cache_interface import FullAttentionSpec
@@ -73,3 +78,85 @@ def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch):
         builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
     )
     assert builder.reorder_batch_threshold == 1
+
+
+def test_dcp_prefill_wrapper_preserves_noncausal_draft_mask() -> None:
+    wrapper = BatchDCPPrefillWrapper.__new__(BatchDCPPrefillWrapper)
+    wrapper._context = Mock()
+    wrapper._new_tokens = Mock()
+    indptr = torch.tensor([0, 3], dtype=torch.int32)
+
+    wrapper.plan(
+        qo_indptr_cpu=indptr,
+        paged_kv_indptr_cpu=torch.tensor([0, 1], dtype=torch.int32),
+        paged_kv_indices=torch.tensor([0], dtype=torch.int32),
+        paged_kv_last_page_len_cpu=torch.tensor([3], dtype=torch.int32),
+        page_size=16,
+        num_qo_heads=4,
+        dcp_world_size=2,
+        num_kv_heads=1,
+        head_dim=128,
+        sm_scale=0.1,
+        window_left=-1,
+        logits_soft_cap=None,
+        q_data_type=torch.bfloat16,
+        kv_cache_dtype=torch.float8_e4m3fn,
+        prefill_fixed_split_size=-1,
+        disable_split_kv=False,
+        causal=False,
+    )
+
+    assert wrapper._context.plan.call_args.kwargs["causal"] is False
+    assert wrapper._new_tokens.plan.call_args.kwargs["causal"] is False
+
+
+def test_flashinfer_selects_dcp_wrapper_for_noncausal_prefill(monkeypatch) -> None:
+    expected = object()
+    factory = Mock(return_value=expected)
+    monkeypatch.setattr(flashinfer_backend, "BatchDCPPrefillWrapper", factory)
+    monkeypatch.setattr(
+        flashinfer_backend,
+        "get_flashinfer_layout_string",
+        lambda _: "NHD",
+    )
+
+    builder = FlashInferMetadataBuilder.__new__(FlashInferMetadataBuilder)
+    builder.use_dcp = True
+    builder.dcp_a2a = False
+    builder._prefill_wrapper = None
+    builder.cache_config = SimpleNamespace(
+        get_resolved_kv_cache_layout=lambda: object()
+    )
+    builder._get_workspace_buffer = Mock(return_value=torch.empty(0))
+
+    assert builder._get_prefill_wrapper(causal=False) is expected
+    factory.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    ("rank", "expected_lengths", "expected_pages"),
+    [
+        (0, [0, 12, 17], [0, 1, 2]),
+        (1, [0, 12, 16], [0, 1, 1]),
+        (2, [0, 12, 16], [0, 1, 1]),
+        (3, [0, 8, 16], [0, 1, 1]),
+    ],
+)
+def test_dcp_flashinfer_page_metadata_is_rank_local(
+    rank: int,
+    expected_lengths: list[int],
+    expected_pages: list[int],
+) -> None:
+    # Global lengths 44 and 65 would incorrectly produce 3 and 5 native
+    # pages at page_size=16. DCP4 stores only the rank-local lengths below.
+    local_lens, local_lens_np, local_pages_np = _get_dcp_local_kv_page_metadata(
+        torch.tensor([0, 44, 65], dtype=torch.int32),
+        dcp_world_size=4,
+        dcp_rank=rank,
+        dcp_kv_cache_interleave_size=4,
+        page_size=16,
+    )
+
+    assert local_lens.tolist() == expected_lengths
+    assert local_lens_np.tolist() == expected_lengths
+    assert local_pages_np.tolist() == expected_pages

tests/v1/core/test_prefix_caching.py

@@ -916,6 +916,81 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
     manager.free(req1)
 
 
+def test_prefill_hybrid_dcp_sliding_window():
+    """DCP scales SWA cache pages, cache-hit lookup, and retention together."""
+    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+
+    block_size = 16
+    dcp_world_size = 4
+    effective_block_size = block_size * dcp_world_size
+    kv_cache_config = _make_hybrid_kv_cache_config(
+        block_size,
+        num_blocks=64,
+        spec_types=["full", "sliding_window_large"],
+    )
+    swa_group = kv_cache_config.kv_cache_groups[1]
+    kv_cache_config = replace(
+        kv_cache_config,
+        kv_cache_groups=[
+            kv_cache_config.kv_cache_groups[0],
+            replace(
+                swa_group,
+                kv_cache_spec=replace(
+                    swa_group.kv_cache_spec,
+                    sliding_window=2 * effective_block_size,
+                ),
+            ),
+        ],
+    )
+    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,
+    )
+
+    assert [m.block_size for m in manager.coordinator.single_type_managers] == [
+        effective_block_size,
+        effective_block_size,
+    ]
+    swa_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec
+    assert isinstance(swa_spec, SlidingWindowSpec)
+    assert SlidingWindowManager.reachable_block_mask(
+        start_block=0,
+        end_block=4,
+        alignment_tokens=effective_block_size,
+        kv_cache_spec=swa_spec,
+        use_eagle=False,
+        retention_interval=0,
+        reachable_boundaries=(3 * effective_block_size - 1,),
+        effective_block_size=effective_block_size,
+    ) == [True, True, False, False]
+
+    common_token_ids = [i for i in range(12) for _ in range(block_size)]
+    req0 = make_request("dcp-producer", common_token_ids + [12] * 7, block_size, sha256)
+    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
+    assert num_computed_tokens == 0
+    assert (
+        manager.allocate_slots(
+            req0,
+            len(req0.prompt_token_ids),
+            num_computed_tokens,
+            computed_blocks,
+        )
+        is not None
+    )
+    manager.new_step_starts()
+
+    req1 = make_request("dcp-consumer", common_token_ids + [13] * 5, block_size, sha256)
+    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
+
+    assert num_computed_tokens == len(common_token_ids)
+    assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 3]
+    manager.free(req0)
+
+
 # Test cases with eagle enabled: Only test a single simple case for now.
 # - 2 groups: 1 full + 1 other
 _EAGLE_HYBRID_MODEL_TEST_CASES = [

tests/v1/core/test_swa_inflight_window_free.py

@@ -168,6 +168,24 @@ def test_swa_admission_cap_accounts_for_overlapping_batches():
     assert overlapped == 193
 
 
+def test_swa_admission_cap_accounts_for_dcp_shards():
+    spec = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        sliding_window=1024,
+    )
+    cap = spec.max_admission_blocks_per_request(
+        max_in_flight_tokens=1024,
+        max_model_len=16384,
+        kv_shard_count=4,
+    )
+    # Each local page covers 16 * DCP4 global tokens. The one extra block
+    # accounts for a window whose trailing edge is not page-aligned.
+    assert cap == 33
+
+
 def test_chunked_local_free_waits_for_in_flight_step():
     """Chunked-local attention frees whole chunks left of the current one, and
     is exposed to the same load-WAR: with async scheduling those chunks must
@@ -221,6 +239,22 @@ def test_chunked_local_admission_cap_accounts_for_overlapping_batches():
     assert overlapped == 192
 
 
+def test_chunked_local_admission_cap_accounts_for_dcp_shards():
+    spec = ChunkedLocalAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        attention_chunk_size=1024,
+    )
+    cap = spec.max_admission_blocks_per_request(
+        max_in_flight_tokens=1024,
+        max_model_len=16384,
+        kv_shard_count=4,
+    )
+    assert cap == 32
+
+
 def test_connector_finish_frees_on_settled_basis():
     """The out-of-window prune done at request finish, before the block table
     is handed to a KV connector (simple CPU offload / NIXL store), must use the