@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 927c87b347

[PCP][DCP] Support DCP target model with non-DCP Dspark (#56723)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Codex <noreply@openai.com>

Lucas Wilkinson committed Sep 26, 2026 at 18:12 UTC 927c87b3472a39da79a95a8df12d0a2b21e2d32e
32 files changed +606 -201
tests/v1/attention/test_attention_backends.py
+1 -2
@@ -942,8 +942,6 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa
942 impl.o_sf_scale = None
943 impl.window_left = -1
944 impl.sinks = None
945 - impl.cache_config = unittest.mock.Mock()
946 - impl.cache_config.get_resolved_kv_cache_layout.return_value = KVCacheLayout.LBHNC
945
946 layer = unittest.mock.Mock(
947 _q_scale=torch.tensor(1.0),
@@ -959,6 +957,7 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa
957 q_len_per_req=1,
958 )
959 attn_metadata = flashinfer_backend.FlashInferMetadata(
960 + kv_cache_layout=KVCacheLayout.LBHNC,
961 num_actual_tokens=2,
962 slot_mapping=torch.empty(0, dtype=torch.int64),
963 q_data_type_prefill=torch.bfloat16,
tests/v1/attention/test_group_head_counts.py
+1
@@ -101,6 +101,7 @@ def test_flash_attention_geometry_comes_from_the_group():
101 num_kv_heads=2,
102 head_size=64,
103 dtype=torch.bfloat16,
104 + dcp_sharded=True,
105 )
106
107 with (
tests/v1/core/test_kv_cache_utils.py
+143 -12
@@ -1599,19 +1599,142 @@ def test_project_kv_cache_groups_to_worker():
1599 assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
1600
1601
1602 -def test_dcp_world_size_for_kv_cache_spec_shards_full_attention_only():
1603 - dcp = 8
1604 - full = FullAttentionSpec(
1605 - block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
1602 +@pytest.mark.parametrize("sliding_window", [None, 256])
1603 +@pytest.mark.parametrize("disable_hybrid", [False, True])
1604 +@pytest.mark.parametrize("pcp_size", [1, 4])
1605 +def test_dcp_target_allocates_replicated_draft_independently(
1606 + monkeypatch, sliding_window, disable_hybrid, pcp_size
1607 +):
1608 + """A draft must retain all positions even when the target shards them."""
1609 + from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
1610 +
1611 + monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
1612 + config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1613 + config.parallel_config.decode_context_parallel_size = 4
1614 + config.parallel_config.prefill_context_parallel_size = pcp_size
1615 + config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hybrid
1616 + config.cache_config.block_size = 16
1617 + config.cache_config.kv_cache_layout = None
1618 + draft_args = dict(
1619 + block_size=16,
1620 + num_kv_heads=8,
1621 + head_size=64,
1622 + dtype=torch.bfloat16,
1623 + dcp_sharded=False,
1624 + )
1625 + draft = (
1626 + FullAttentionSpec(**draft_args)
1627 + if sliding_window is None
1628 + else SlidingWindowSpec(**draft_args, sliding_window=sliding_window)
1629 )
1607 - mla = new_mla_spec()
1608 - mamba = new_mamba_spec()
1609 - uniform_mla = UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs={"layer": mla})
1610 - assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(full, dcp) == dcp
1611 - assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(mla, dcp) == dcp
1612 - assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(uniform_mla, dcp) == dcp
1613 - assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(mamba, dcp) == 1
1614 - assert kv_cache_utils.dcp_world_size_for_kv_cache_spec(full, 1) == 1
1630 + specs = {"target": new_mla_spec(), "draft": draft}
1631 + layout = resolve_kv_cache_layout(config, [["LBHNC", "BLHNC"]], specs.values())
1632 + assert layout == KVCacheLayout.BLHNC
1633 + groups = get_kv_cache_groups(config, specs)
1634 + assert [group.layer_names for group in groups] == [["target"], ["draft"]]
1635 + widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups]
1636 + assert widths == [16, 64]
1637 +
1638 + cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
1639 + config, groups, available_memory=16 * 1024 * 1024
1640 + )
1641 + scheduler_config = generate_scheduler_kv_cache_config([cache_config])
1642 + assert get_max_concurrency_for_kv_cache_config(config, cache_config) > 0
1643 + assert get_max_concurrency_for_kv_cache_config(config, cache_config) == (
1644 + get_max_concurrency_for_kv_cache_config(config, scheduler_config)
1645 + )
1646 + manager = KVCacheManager(
1647 + scheduler_config,
1648 + max_model_len=1024,
1649 + hash_block_size=16,
1650 + scheduler_block_size=64,
1651 + dcp_world_size=4,
1652 + pcp_world_size=pcp_size,
1653 + enable_caching=True,
1654 + )
1655 + assert manager.coordinator.group_block_sizes == (64, 16)
1656 + request = make_request("replicated-draft", [1] * 65, block_size=16, hash_fn=sha256)
1657 + blocks = manager.allocate_slots(request, 65)
1658 + assert blocks is not None
1659 + assert [len(group) for group in blocks.blocks] == [2, 5]
1660 + manager.cache_blocks(request, 64)
1661 + cached_request = make_request(
1662 + "cached-draft", [1] * 65, block_size=16, hash_fn=sha256
1663 + )
1664 + cached_blocks, num_cached, _ = manager.get_computed_blocks(cached_request)
1665 + assert num_cached == 64
1666 + assert [len(group) for group in cached_blocks.blocks] == [1, 4]
1667 +
1668 +
1669 +@pytest.mark.parametrize("use_mla", [False, True])
1670 +def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla):
1671 + config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1672 + config.parallel_config.decode_context_parallel_size = 4
1673 + draft = replace(
1674 + new_mla_spec() if use_mla else new_kv_cache_spec(),
1675 + block_size=16,
1676 + dcp_sharded=False,
1677 + )
1678 + merged = type(draft).merge([draft, draft])
1679 + wrapped = UniformTypeKVCacheSpecs.from_specs({"draft": merged})
1680 + assert wrapped is not None and not wrapped.dcp_sharded
1681 + assert merged.max_num_blocks_per_req(config, 1024) == 64
1682 + assert merged.max_memory_usage_bytes(config) == 64 * draft.page_size_bytes
1683 + assert kv_cache_utils.resolve_dcp_kv_block_size(merged, 4) == 16
1684 + with pytest.raises(AssertionError):
1685 + type(draft).merge([draft, replace(draft, dcp_sharded=True)])
1686 +
1687 +
1688 +@pytest.mark.parametrize("with_draft", [False, True])
1689 +@pytest.mark.parametrize("indexer_alignment", [None, 512])
1690 +def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment):
1691 + from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view
1692 + from vllm.v1.worker.utils import allocate_kv_cache
1693 +
1694 + config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1695 + config.cache_config.kv_cache_layout = "BLHNC"
1696 + config.parallel_config.decode_context_parallel_size = 4
1697 + common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8)
1698 + specs = {
1699 + "target": MLAAttentionSpec(
1700 + **common,
1701 + head_size=576,
1702 + state_content_bytes=656,
1703 + block_stride_alignment=656,
1704 + cache_dtype_str="fp8_ds_mla",
1705 + is_index_group_leader=True,
1706 + ),
1707 + "indexer": MLAAttentionSpec(
1708 + **common,
1709 + head_size=132,
1710 + cache_role=SparseCacheRole.INDEXER,
1711 + block_stride_alignment=indexer_alignment,
1712 + ),
1713 + "draft": SlidingWindowSpec(
1714 + block_size=64,
1715 + num_kv_heads=64,
1716 + head_size=64,
1717 + dtype=torch.bfloat16,
1718 + sliding_window=2048,
1719 + dcp_sharded=False,
1720 + ),
1721 + }
1722 + if not with_draft:
1723 + del specs["draft"]
1724 + groups = get_kv_cache_groups(config, specs)
1725 + cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
1726 + config, groups, available_memory=8 * 1024 * 1024
1727 + )
1728 + caches = allocate_kv_cache(
1729 + cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups)
1730 + )
1731 + cache = caches["target"].squeeze(1)
1732 + if indexer_alignment is not None:
1733 + assert caches["indexer"].stride(0) % indexer_alignment == 0
1734 + rows, stride_rows = flat_kv_row_view(cache, 64)
1735 + cache[1, 0].fill_(7)
1736 + torch.testing.assert_close(rows[stride_rows], cache[1, 0])
1737 + assert torch.all(rows[stride_rows] == 7)
1738
1739
1740 @pytest.mark.parametrize(
@@ -3815,6 +3938,8 @@ def test_unify_hybrid_kv_cache_specs():
3938 "layer_1": before_spec_1,
3939 "layer_2": before_spec_2,
3940 }
3941 + kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False)
3942 + kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False)
3943 kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
3944 expected_spec_1 = new_kv_cache_spec(block_size=64)
3945 expected_spec_2 = new_kv_cache_spec(
@@ -3823,6 +3948,7 @@ def test_unify_hybrid_kv_cache_specs():
3948 assert kv_cache_spec["layer_1"] == expected_spec_1
3949 assert kv_cache_spec["layer_2"] == expected_spec_2
3950 assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024
3951 + assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False)
3952
3953 # 2. has_full_attention and has_chunked_local_attention
3954 before_spec_1 = new_kv_cache_spec()
@@ -3874,6 +4000,11 @@ def test_unify_hybrid_kv_cache_specs():
4000 with pytest.raises(ValueError):
4001 kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4002
4003 + # Replicated Mamba state still requires the hybrid cache manager.
4004 + kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()}
4005 + with pytest.raises(ValueError):
4006 + kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
4007 +
4008
4009 def test_unify_kv_cache_spec_page_size_mamba():
4010 """Regression test for https://github.com/vllm-project/vllm/issues/43626.
tests/v1/core/test_prefix_caching.py
+12 -7
@@ -1336,12 +1336,12 @@ def make_kv_cache_config_three_types(
1336 )
1337
1338
1339 -def test_prefix_cache_hit_uses_per_group_dcp_geometry():
1339 +@pytest.mark.parametrize("draft_sharded", [True, False])
1340 +def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded):
1341 """Prefix lookup must use each group's DCP size, not the process-wide one.
1342
1342 - Target and draft MLA are both sharded (DCP=8); Mamba stays replicated
1343 - (DCP=1). Hits then align to the sharded full-attention block, not the
1344 - unsharded page size.
1343 + Target MLA is sharded; draft MLA can opt out and Mamba stays replicated.
1344 + Hits must align to the target block in either case.
1345 """
1346 block_size = 16
1347 dcp = 8
@@ -1362,6 +1362,7 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry():
1362 KVCacheGroupSpec(
1363 ["draft_mla"],
1364 MLAAttentionSpec(
1365 + dcp_sharded=draft_sharded,
1366 block_size=block_size,
1367 num_kv_heads=1,
1368 head_size=1,
@@ -1388,10 +1389,10 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry():
1389 )
1390 target_mgr, draft_mgr, mamba_mgr = manager.coordinator.single_type_managers
1391 assert target_mgr.dcp_world_size == dcp
1391 - assert draft_mgr.dcp_world_size == dcp
1392 + assert draft_mgr.dcp_world_size == (dcp if draft_sharded else 1)
1393 assert mamba_mgr.dcp_world_size == 1
1394 assert target_mgr.block_size == sharded_block
1394 - assert draft_mgr.block_size == sharded_block
1395 + assert draft_mgr.block_size == (sharded_block if draft_sharded else block_size)
1396 assert mamba_mgr.block_size == block_size
1397
1398 hash_fn = sha256
@@ -1407,7 +1408,11 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry():
1408 req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn)
1409 computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
1410 assert num_computed_tokens == 2 * sharded_block
1410 - assert [len(group) for group in computed_blocks.blocks] == [2, 2, 16]
1411 + assert [len(group) for group in computed_blocks.blocks] == [
1412 + 2,
1413 + 2 if draft_sharded else 16,
1414 + 16,
1415 + ]
1416
1417 manager.free(req0)
1418 manager.free(req1)
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
+33
@@ -11,6 +11,7 @@ from vllm.config import CompilationConfig, VllmConfig, replace
11 from vllm.config.kernel import MoEBackend
12 from vllm.engine.arg_utils import EngineArgs
13 from vllm.platforms import current_platform
14 +from vllm.v1.spec_decode.draft_model import DraftModelProposer
15
16 from ..utils import (
17 Messages,
@@ -196,6 +197,38 @@ def test_draft_model_engine_args_tensor_parallelism():
197 assert draft_config.quant_config is None
198
199
200 +@pytest.mark.cpu_test
201 +@pytest.mark.parametrize("dcp_size", [1, 2])
202 +def test_draft_model_inherits_target_dcp_config(monkeypatch, dcp_size):
203 + monkeypatch.setattr(current_platform, "device_count", lambda: 16)
204 + target_config = EngineArgs(
205 + model="Qwen/Qwen3-1.7B",
206 + tensor_parallel_size=16,
207 + decode_context_parallel_size=dcp_size,
208 + cp_kv_cache_interleave_size=16,
209 + dcp_comm_backend="ag_rs",
210 + distributed_executor_backend="mp",
211 + speculative_config={
212 + "model": "Qwen/Qwen3-0.6B",
213 + "method": "draft_model",
214 + "num_speculative_tokens": 3,
215 + },
216 + ).create_engine_config()
217 + proposer = object.__new__(DraftModelProposer)
218 + proposer.vllm_config = target_config
219 + proposer.speculative_config = target_config.speculative_config
220 + draft_config = proposer._create_draft_vllm_config()
221 + draft_parallel = draft_config.parallel_config
222 + assert draft_parallel.decode_context_parallel_size == dcp_size
223 + assert draft_parallel.cp_kv_cache_interleave_size == 16
224 + assert draft_parallel.dcp_comm_backend == "ag_rs"
225 + assert (
226 + draft_parallel.dcp_q_replicate == target_config.parallel_config.dcp_q_replicate
227 + )
228 + assert draft_parallel is not target_config.parallel_config
229 + assert target_config.parallel_config.decode_context_parallel_size == dcp_size
230 +
231 +
232 def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
233 """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic
234 so we can test it without instantiating a full proposer."""
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+2 -1
@@ -500,7 +500,7 @@ def test_decode_bench_connector_skips_null_padding():
500
501
502 def test_decode_bench_connector_uses_per_group_dcp_block_sizes():
503 - """DCP scales full-attention blocks but not sliding-window blocks."""
503 + """DCP scales target blocks but not the replicated draft's window."""
504 block_size = 16
505 dcp_world_size = 2
506 vllm_config = create_vllm_config(
@@ -515,6 +515,7 @@ def test_decode_bench_connector_uses_per_group_dcp_block_sizes():
515 dtype=torch.float32,
516 )
517 sliding_window = SlidingWindowSpec(
518 + dcp_sharded=False,
519 block_size=block_size,
520 num_kv_heads=1,
521 head_size=1,
tests/v1/simple_kv_offload/test_scheduler.py
+2 -2
@@ -2316,8 +2316,8 @@ def _make_hybrid_attention_mamba_scheduler(
2316 def test_hybrid_store_uses_resolved_group_block_sizes() -> None:
2317 """Replicated groups must not be scaled by the DCP world size.
2318
2319 - ``dcp_world_size_for_kv_cache_spec`` gives full attention the process DCP
2320 - size and every other spec 1. Reconstructing a group's geometry as
2319 + The spec's ``dcp_sharded`` flag controls its effective block size.
2320 + Reconstructing a group's geometry as
2321 ``spec.block_size * cp_world_size`` over-scales the replicated groups, so
2322 the eager store scan believes far fewer of their blocks are ready and
2323 silently offloads only a fraction of them.
tests/v1/spec_decode/test_dflash_prepare_inputs.py
+68 -30
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5 +from unittest.mock import patch
6
7 import numpy as np
8 import pytest
@@ -9,7 +10,7 @@ import torch
10
11 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
12 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
12 - prepare_dflash_inputs,
13 + DFlashSpeculator,
14 )
15
16 pytestmark = pytest.mark.skipif(
@@ -24,6 +25,7 @@ def _run_prepare(
25 cp_rank: int = 0,
26 cp_size: int = 1,
27 cp_interleave: int = 1,
28 + draft_dcp_size: int | None = None,
29 ):
30 device = torch.device("cuda")
31 max_num_reqs = 4
@@ -40,7 +42,9 @@ def _run_prepare(
42 )
43 input_batch = SimpleNamespace(
44 num_reqs=1,
45 + num_tokens=4,
46 num_scheduled_tokens=np.array([4], dtype=np.int32),
47 + seq_lens_cpu_upper_bound=torch.tensor([target_positions[-1] + 1]),
48 positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
49 query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
50 idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),
@@ -74,36 +78,65 @@ def _run_prepare(
78 next_prefill_tokens = torch.zeros_like(last_sampled)
79 block_table = torch.tensor([block_table_values], dtype=torch.int32, device=device)
80
77 - prepare_dflash_inputs(
78 - input_buffers,
79 - query_slot_mapping,
80 - context_positions,
81 - context_slot_mapping,
82 - sample_indices,
83 - sample_pos,
84 - sample_idx_mapping,
85 - temperature,
86 - seeds,
87 - input_batch,
88 - torch.tensor([1], dtype=torch.int32, device=device),
89 - torch.tensor([2], dtype=torch.int32, device=device),
90 - last_sampled,
91 - next_prefill_tokens,
92 - input_temperature,
93 - input_seeds,
94 - block_table,
95 - 4,
96 - cp_rank,
97 - cp_size,
98 - cp_interleave,
99 - 123,
100 - num_speculative_steps,
101 - num_speculative_steps,
102 - max_num_reqs,
103 - max_num_tokens,
104 - 128,
81 + class InputsPrepared(Exception):
82 + pass
83 +
84 + draft = SimpleNamespace(
85 + input_buffers=input_buffers,
86 + context_positions=context_positions,
87 + _context_slot_mappings=context_slot_mapping.unsqueeze(0),
88 + sample_indices=sample_indices,
89 + sample_pos=sample_pos,
90 + sample_idx_mapping=sample_idx_mapping,
91 + temperature=temperature,
92 + seeds=seeds,
93 + hidden_states=torch.zeros(4, 1, device=device),
94 + prepare_context_anchor=lambda *args: None,
95 + query_cudagraph_manager=None,
96 + dp_size=1,
97 + dp_rank=0,
98 + pcp_manager=None,
99 + draft_kv_cache_group_id=0,
100 + draft_kv_cache_group_ids=[0],
101 + _layer_group_idx=None,
102 + block_tables=SimpleNamespace(
103 + slot_mappings=query_slot_mapping.unsqueeze(0),
104 + input_block_tables=[block_table],
105 + kernel_block_sizes=[4],
106 + cp_rank=cp_rank,
107 + cp_size=cp_size,
108 + cp_interleave=cp_interleave,
109 + ),
110 + dcp_size=cp_size if draft_dcp_size is None else draft_dcp_size,
111 + parallel_drafting_token_id=123,
112 + num_query_per_req=num_speculative_steps,
113 + num_speculative_steps=num_speculative_steps,
114 + max_num_reqs=max_num_reqs,
115 + max_num_tokens=max_num_tokens,
116 + max_model_len=128,
117 sample_from_anchor=True,
118 )
119 + with (
120 + patch(
121 + "vllm.v1.worker.gpu.spec_decode.dflash.speculator.dispatch_cg_and_sync_dp",
122 + side_effect=InputsPrepared,
123 + ),
124 + pytest.raises(InputsPrepared),
125 + ):
126 + DFlashSpeculator.propose(
127 + draft,
128 + input_batch,
129 + {},
130 + {},
131 + draft.hidden_states,
132 + None,
133 + torch.tensor([1], dtype=torch.int32, device=device),
134 + torch.tensor([2], dtype=torch.int32, device=device),
135 + last_sampled,
136 + next_prefill_tokens,
137 + input_temperature,
138 + input_seeds,
139 + )
140 torch.accelerator.synchronize()
141 return SimpleNamespace(
142 input_buffers=input_buffers,
@@ -118,12 +151,17 @@ def _run_prepare(
151 )
152
153
121 -def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
154 +@pytest.mark.parametrize("cp_rank,cp_size", [(0, 1), (0, 2), (1, 2)])
155 +def test_prepare_dflash_inputs_excludes_rejected_context_suffix(cp_rank, cp_size):
156 # Positions 10/11 use physical block 7. Rejected positions 12/13 would use
157 # block 8, but must be PAD context rather than contaminating draft KV.
158 out = _run_prepare(
159 target_positions=[10, 11, 12, 13],
160 block_table_values=[0, 0, 7, 8, 9, 10, 11, 12],
161 + cp_rank=cp_rank,
162 + cp_size=cp_size,
163 + cp_interleave=2,
164 + draft_dcp_size=1,
165 )
166
167 assert out.context_positions[:4].tolist() == [10, 11, 0, 0]
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+9 -1
@@ -17,6 +17,7 @@ from types import MethodType, SimpleNamespace
17 from unittest.mock import patch
18
19 import numpy as np
20 +import pytest
21 import torch
22
23 from vllm.config.compilation import CUDAGraphMode
@@ -53,6 +54,7 @@ def _make_fake_speculator(
54 arange_np=np.arange(max_num_reqs + 1, dtype=np.int32),
55 draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool),
56 block_tables=fake_block_tables,
57 + dcp_size=1,
58 input_buffers=fake_input_buffers,
59 attn_groups=[],
60 kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
@@ -152,14 +154,17 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len():
154 assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
155
156
155 -def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
157 +@pytest.mark.parametrize("draft_dcp_size", [1, 2])
158 +def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens(draft_dcp_size):
159 fake = _make_fake_speculator()
160 + fake.dcp_size = draft_dcp_size
161 fake.block_tables.cp_size = 2
162 fake.block_tables.cp_rank = 1
163 fake.block_tables.cp_interleave = 4
164 fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16])
165
166 def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave):
167 + assert draft_dcp_size > 1
168 assert seq_lens is fake.input_buffers.seq_lens
169 assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4)
170 out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32))
@@ -176,6 +181,9 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
181 )
182
183 local = captured["dcp_local_seq_lens"]
184 + if draft_dcp_size == 1:
185 + assert local is None
186 + return
187 assert isinstance(local, torch.Tensor)
188 assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr()
189 assert local.tolist() == [1, 4, 8, 0]
tests/v1/worker/test_cp_utils.py
+34
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from types import SimpleNamespace
4 +
5 import pytest
6 import torch
7
@@ -7,6 +9,38 @@ from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
9 from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention
10
11
12 +def test_cp_checks_allow_a_replicated_draft_backend_without_dcp(monkeypatch):
13 + from vllm.v1.worker import cp_utils
14 +
15 + config = SimpleNamespace(
16 + parallel_config=SimpleNamespace(
17 + prefill_context_parallel_size=1,
18 + decode_context_parallel_size=4,
19 + cp_kv_cache_interleave_size=16,
20 + ),
21 + speculative_config=object(),
22 + )
23 + target = SimpleNamespace(
24 + impl=SimpleNamespace(
25 + dcp_world_size=4,
26 + supports_mtp_with_cp_non_trivial_interleave_size=True,
27 + need_to_return_lse_for_decode=True,
28 + )
29 + )
30 + draft = SimpleNamespace(
31 + impl=SimpleNamespace(
32 + dcp_world_size=1,
33 + supports_mtp_with_cp_non_trivial_interleave_size=False,
34 + need_to_return_lse_for_decode=False,
35 + )
36 + )
37 + layers = {"target": target, "draft": draft}
38 + monkeypatch.setattr(cp_utils, "get_layers_from_vllm_config", lambda *args: layers)
39 + cp_utils.check_attention_cp_compatibility(config, target_layer_names={"target"})
40 + with pytest.raises(AssertionError, match="MTP with cp_kv_cache_interleave_size"):
41 + cp_utils.check_attention_cp_compatibility(config)
42 +
43 +
44 def test_skip_gate_only_for_zero_context():
45 assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32))
46 assert not should_skip_dcp_context_attention(
tests/v1/worker/test_dspark_utils.py
+92
@@ -2,8 +2,17 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass, field
5 +from types import SimpleNamespace
6
7 +import pytest
8 +import torch
9 +
10 +from vllm.config import ParallelConfig
11 +from vllm.v1.kv_cache_interface import FullAttentionSpec
12 +from vllm.v1.worker.gpu import model_runner
13 +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
14 from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config
15 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
16
17
18 @dataclass
@@ -49,3 +58,86 @@ def test_dspark_parallel_config_disables_eplb_atomically():
58 assert draft_config.eplb_config.num_redundant_experts == 0
59 assert not draft_config.enable_elastic_ep
60 assert draft_config.eplb_config is not target_config.eplb_config
61 +
62 +
63 +@pytest.mark.parametrize("pcp_size", [1, 4])
64 +@pytest.mark.parametrize("dcp_size", [1, 4])
65 +@pytest.mark.parametrize("use_mla", [False, True])
66 +def test_draft_context_parallelism_without_changing_target(
67 + monkeypatch, pcp_size, dcp_size, use_mla
68 +):
69 + target_parallel = ParallelConfig(
70 + tensor_parallel_size=4,
71 + prefill_context_parallel_size=pcp_size,
72 + decode_context_parallel_size=dcp_size,
73 + cp_kv_cache_interleave_size=16,
74 + distributed_executor_backend="mp",
75 + )
76 + target_config = SimpleNamespace(
77 + parallel_config=target_parallel,
78 + speculative_config=SimpleNamespace(
79 + draft_model_config=SimpleNamespace(use_mla=use_mla)
80 + ),
81 + )
82 +
83 + class CapturedConfig(Exception):
84 + pass
85 +
86 + def capture_init(self, config, device):
87 + raise CapturedConfig(config)
88 +
89 + monkeypatch.setattr(DraftModelSpeculator, "__init__", capture_init)
90 + with pytest.raises(CapturedConfig) as captured:
91 + DFlashSpeculator(target_config, device=None)
92 + draft_parallel = captured.value.args[0].parallel_config
93 + assert draft_parallel.tensor_parallel_size == 4
94 + assert draft_parallel.prefill_context_parallel_size == 1
95 + assert draft_parallel.decode_context_parallel_size == (dcp_size if use_mla else 1)
96 + assert draft_parallel.cp_kv_cache_interleave_size == 16
97 + assert target_parallel.prefill_context_parallel_size == pcp_size
98 + assert target_parallel.decode_context_parallel_size == dcp_size
99 + assert target_parallel.cp_kv_cache_interleave_size == 16
100 +
101 +
102 +@pytest.mark.parametrize("configured_dcp", [1, 4])
103 +def test_attention_uses_draft_dcp_setting_inside_target_process_group(
104 + monkeypatch, configured_dcp
105 +):
106 + import vllm.config as config_module
107 + from vllm.distributed import parallel_state
108 + from vllm.v1.attention.backend import AttentionImplBase
109 +
110 + config = SimpleNamespace(
111 + parallel_config=SimpleNamespace(decode_context_parallel_size=configured_dcp)
112 + )
113 + monkeypatch.setattr(
114 + config_module, "get_current_vllm_config_or_none", lambda: config
115 + )
116 +
117 + monkeypatch.setattr(
118 + parallel_state, "_DCP", SimpleNamespace(world_size=4, rank_in_group=2)
119 + )
120 + impl = AttentionImplBase()
121 + assert impl.dcp_world_size == configured_dcp
122 + assert impl.dcp_rank == (0 if configured_dcp == 1 else 2)
123 +
124 +
125 +@pytest.mark.parametrize("draft_dcp_size", [1, 4])
126 +def test_runner_marks_only_replicated_draft_caches(monkeypatch, draft_dcp_size):
127 + spec = FullAttentionSpec(
128 + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.bfloat16
129 + )
130 + monkeypatch.setattr(
131 + model_runner, "get_kv_cache_spec", lambda _: {"target": spec, "draft": spec}
132 + )
133 + runner = object.__new__(model_runner.GPUModelRunner)
134 + runner.vllm_config = None
135 + runner.dcp_size = 4
136 + runner.speculator = object.__new__(DFlashSpeculator)
137 + runner.speculator.dcp_size = draft_dcp_size
138 + runner.speculator.draft_attn_layer_names = {"draft"}
139 +
140 + specs = runner.get_kv_cache_spec()
141 +
142 + assert specs["target"].dcp_sharded
143 + assert specs["draft"].dcp_sharded == (draft_dcp_size == 4)
tests/v1/worker/test_gpu_autoregressive_speculator.py
+2
@@ -88,6 +88,7 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):
88 query_start_loc=torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32),
89 seq_lens=torch.tensor([11, 21, 0, 0], dtype=torch.int32),
90 )
91 + speculator.dcp_size = 1
92 speculator.block_tables = SimpleNamespace(
93 cp_size=1,
94 input_block_tables=[torch.zeros(4, 1, dtype=torch.int32)],
@@ -183,6 +184,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected)
184 parallel_config=SimpleNamespace(
185 data_parallel_size=1,
186 data_parallel_rank=0,
187 + decode_context_parallel_size=1,
188 ),
189 )
190
tests/v1/worker/test_gpu_block_table.py
+9 -7
@@ -171,22 +171,23 @@ def test_block_tables_skip_custom_slot_mapping_groups():
171
172 @pytest.mark.parametrize("cp_rank", range(4))
173 def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
174 - """DCP interleave is expressed in logical-block token coordinates."""
174 + """Only sharded groups use DCP interleave in logical-block coordinates."""
175 device = torch.device("cuda")
176 block_tables = BlockTables(
177 - block_sizes=[128],
177 + block_sizes=[128, 128],
178 max_num_reqs=1,
179 max_num_batched_tokens=1024,
180 - max_num_blocks_per_group=[2],
180 + max_num_blocks_per_group=[2, 8],
181 device=device,
182 - kernel_block_sizes=[64],
182 + kernel_block_sizes=[64, 64],
183 + dcp_sharded=[True, False],
184 cp_size=4,
185 cp_rank=cp_rank,
186 cp_interleave=128,
187 )
188 block_tables.append_block_ids(
189 req_index=0,
189 - new_block_ids=([5, 9],),
190 + new_block_ids=([5, 9], list(range(10, 18))),
191 overwrite=True,
192 )
193 block_tables.apply_staged_writes()
@@ -199,7 +200,7 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
200 query_start_loc,
201 positions,
202 num_tokens_padded=1024,
202 - )[0]
203 + )
204
205 expected = torch.full((1024,), -1, dtype=torch.int64, device=device)
206 first_start = cp_rank * 128
@@ -210,7 +211,8 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
211 expected[second_start : second_start + 128] = torch.arange(
212 9 * 128, 10 * 128, dtype=torch.int64, device=device
213 )
213 - assert torch.equal(actual, expected)
214 + assert torch.equal(actual[0], expected)
215 + assert torch.equal(actual[1], positions + 10 * 128)
216
217
218 def test_v1_block_table_move_row_clears_vacated_row():
tests/v1/worker/test_gpu_pcp_manager.py
+18
@@ -75,6 +75,24 @@ def test_validate_config_rejects_full_graph_for_prefills():
75 )
76
77
78 +@pytest.mark.parametrize("draft_uses_mla", [False, True])
79 +@pytest.mark.parametrize("dcp_size", [1, 2, 4])
80 +def test_validate_config_dcp_for_dspark(draft_uses_mla, dcp_size):
81 + config = _make_config(CUDAGraphMode.NONE)
82 + config.parallel_config.prefill_context_parallel_size = 4
83 + config.parallel_config.decode_context_parallel_size = dcp_size
84 + config.model_config.hf_text_config.index_topk = 1
85 + config.speculative_config = SimpleNamespace(
86 + use_dspark=lambda: True,
87 + draft_model_config=SimpleNamespace(use_mla=draft_uses_mla),
88 + )
89 + if draft_uses_mla and dcp_size == 2:
90 + with pytest.raises(NotImplementedError, match="requires DCP=1 or DCP=PCP"):
91 + PCPManager.validate_config(config, supports_mm_inputs=False)
92 + return
93 + PCPManager.validate_config(config, supports_mm_inputs=False)
94 +
95 +
96 def test_replicated_decode_piecewise_graph_padding(monkeypatch):
97 manager = PCPManager(
98 pcp_world_size=2,
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+1 -8
@@ -59,7 +59,6 @@ from vllm.platforms import current_platform
59 from vllm.utils.math_utils import cdiv
60 from vllm.v1.attention.backend import AttentionMetadata
61 from vllm.v1.core.kv_cache_utils import (
62 - dcp_world_size_for_kv_cache_spec,
62 resolve_dcp_kv_block_size,
63 )
64 from vllm.v1.kv_cache_interface import (
@@ -259,13 +258,7 @@ class DecodeBenchConnectorScheduler:
258 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
259 self.kv_cache_groups = kv_cache_config.kv_cache_groups
260 self.group_block_sizes = tuple(
262 - resolve_dcp_kv_block_size(
263 - group.kv_cache_spec,
264 - dcp_world_size_for_kv_cache_spec(
265 - group.kv_cache_spec,
266 - dcp_world_size,
267 - ),
268 - )
261 + resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_world_size)
262 for group in self.kv_cache_groups
263 )
264
vllm/distributed/parallel_state.py
+8
@@ -1583,6 +1583,14 @@ def get_dcp_group() -> GroupCoordinator:
1583 return _DCP
1584
1585
1586 +def get_dcp_world_size_and_rank(enabled: bool = True) -> tuple[int, int]:
1587 + """Return ``(world_size, rank)`` in the DCP group, or ``(1, 0)`` when disabled
1588 + (e.g. a replicated draft cache) or the group is uninitialized (unit tests)."""
1589 + if not enabled or _DCP is None:
1590 + return 1, 0
1591 + return _DCP.world_size, _DCP.rank_in_group
1592 +
1593 +
1594 _PP: GroupCoordinator | None = None
1595
1596
vllm/v1/attention/backend.py
+7 -8
@@ -870,15 +870,14 @@ class AttentionImplBase(ABC, Generic[T]):
870 def __new__(cls, *args, **kwargs):
871 # use __new__ so that all subclasses will call this
872 self = super().__new__(cls)
873 - try:
874 - from vllm.distributed.parallel_state import get_dcp_group
873 + from vllm.config import get_current_vllm_config_or_none
874 + from vllm.distributed.parallel_state import get_dcp_world_size_and_rank
875
876 - self.dcp_world_size = get_dcp_group().world_size
877 - self.dcp_rank = get_dcp_group().rank_in_group
878 - except AssertionError:
879 - # DCP might not be initialized in testing
880 - self.dcp_world_size = 1
881 - self.dcp_rank = 0
876 + # Replicated drafts run at DCP=1 inside a DCP target's process group.
877 + config = get_current_vllm_config_or_none()
878 + self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
879 + config is None or config.parallel_config.decode_context_parallel_size > 1
880 + )
881 try:
882 from vllm.distributed.parallel_state import get_pcp_group
883
vllm/v1/attention/backends/flash_attn.py
+7 -10
@@ -60,7 +60,10 @@ from vllm.config import (
60 get_layers_from_vllm_config,
61 )
62 from vllm.config.cache import CacheDType
63 -from vllm.distributed.parallel_state import get_dcp_group
63 +from vllm.distributed.parallel_state import (
64 + get_dcp_group,
65 + get_dcp_world_size_and_rank,
66 +)
67 from vllm.logger import init_logger
68 from vllm.platforms.interface import DeviceCapability
69 from vllm.utils.math_utils import cdiv, round_up
@@ -726,15 +729,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
729 num_qo_heads=self.num_heads_q,
730 )
731
729 - try:
730 - from vllm.distributed.parallel_state import get_dcp_group
731 -
732 - self.dcp_world_size = get_dcp_group().world_size
733 - self.dcp_rank = get_dcp_group().rank_in_group
734 - except AssertionError:
735 - # DCP might not be initialized in testing
736 - self.dcp_world_size = 1
737 - self.dcp_rank = 0
732 + self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
733 + kv_cache_spec.dcp_sharded
734 + )
735
736 # Fused draft decode reuses the captured metadata object across draft
737 # steps. For DCP, build-time host-side decisions such as
vllm/v1/attention/backends/flashinfer.py
+19 -25
@@ -30,7 +30,10 @@ from vllm.config import (
30 get_current_vllm_config_or_none,
31 )
32 from vllm.config.cache import CacheDType
33 -from vllm.distributed.parallel_state import get_dcp_group
33 +from vllm.distributed.parallel_state import (
34 + get_dcp_group,
35 + get_dcp_world_size_and_rank,
36 +)
37 from vllm.logger import init_logger
38 from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 QuantKey,
@@ -641,6 +644,7 @@ class FlashInferTrtllmAPIDecode:
644
645 @dataclass
646 class FlashInferMetadata:
647 + kv_cache_layout: KVCacheLayout
648 num_actual_tokens: int
649 """Total number of tokens in the batch (excluding padding)."""
650
@@ -747,17 +751,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
751 self._decode_cudagraph_max_bs,
752 self.compilation_config.max_cudagraph_capture_size,
753 )
750 - try:
751 - self.dcp_world_size = get_dcp_group().world_size
752 - self.dcp_rank = get_dcp_group().rank_in_group
753 - self.dcp_kv_cache_interleave_size = (
754 - vllm_config.parallel_config.cp_kv_cache_interleave_size
755 - )
756 - except AssertionError:
757 - # DCP might not be initialized in testing
758 - self.dcp_world_size = 1
759 - self.dcp_rank = 0
760 - self.dcp_kv_cache_interleave_size = 1
754 + self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
755 + kv_cache_spec.dcp_sharded
756 + )
757 + self.dcp_kv_cache_interleave_size = (
758 + vllm_config.parallel_config.cp_kv_cache_interleave_size
759 + )
760 self.use_dcp = self.dcp_world_size > 1
761 self.dcp_a2a = (
762 self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
@@ -1424,6 +1423,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1423 # Leave prefill/decode/cascade_wrapper empty, to be populated
1424 # case by case depending on the batch contents and backend selection.
1425 attn_metadata = FlashInferMetadata(
1426 + kv_cache_layout=self.kv_cache_layout,
1427 num_actual_tokens=num_actual_tokens,
1428 slot_mapping=common_attn_metadata.slot_mapping,
1429 q_data_type_prefill=self.q_data_type_prefill,
@@ -1904,8 +1904,6 @@ class FlashInferImpl(AttentionImpl):
1904 num_heads, num_kv_heads, is_prefill=False
1905 )
1906 vllm_config = get_current_vllm_config_or_none()
1907 - # The layout is resolved after model construction, so read it lazily.
1908 - self.cache_config = vllm_config.cache_config if vllm_config else None
1907 # Query pre-quantization needs a single dtype for the whole query tensor.
1908 # SM90 XQA needs BF16/FP16-Q for decode and FP8 for prefill,
1909 # so only enable this for SM100 trtllm-gen where both use FP8-Q.
@@ -1941,11 +1939,6 @@ class FlashInferImpl(AttentionImpl):
1939 else:
1940 self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
1941
1944 - @property
1945 - def kv_cache_layout(self) -> KVCacheLayout:
1946 - assert self.cache_config is not None
1947 - return self.cache_config.get_resolved_kv_cache_layout()
1948 -
1942 def fused_output_quant_supported(self, quant_key: QuantKey):
1943 if quant_key == kNvfp4Dynamic and self.is_kvcache_nvfp4:
1944 logger.warning_once(
@@ -2111,6 +2104,7 @@ class FlashInferImpl(AttentionImpl):
2104 # performance to make sure it does not introduce any overhead.
2105
2106 num_actual_tokens = attn_metadata.num_actual_tokens
2107 + kv_cache_layout = attn_metadata.kv_cache_layout
2108
2109 # FlashInfer treats uint8 KV cache as NVFP4. vLLM stores FP8 KV cache
2110 # as uint8 bytes, so pass FP8 caches with their logical dtype.
@@ -2134,7 +2128,7 @@ class FlashInferImpl(AttentionImpl):
2128 if attn_metadata.use_cascade:
2129 # Cascade attention (rare case).
2130 assert attn_metadata.cascade_wrapper is not None
2137 - stride_order = self.kv_cache_layout.layer_view_order
2131 + stride_order = kv_cache_layout.layer_view_order
2132 if self.is_kvcache_nvfp4:
2133 kv_cache_views = tuple(
2134 cache.permute(*stride_order)
@@ -2154,7 +2148,7 @@ class FlashInferImpl(AttentionImpl):
2148 num_decode_tokens = attn_metadata.num_decode_tokens
2149 num_prefill_tokens = attn_metadata.num_prefill_tokens
2150
2157 - stride_order = self.kv_cache_layout.layer_view_order
2151 + stride_order = kv_cache_layout.layer_view_order
2152 kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
2153 # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1
2154 # with TP=8). PyTorch permits non-canonical strides on size-1 dims;
@@ -2311,7 +2305,7 @@ class FlashInferImpl(AttentionImpl):
2305 seq_lens_prefill = attn_metadata.prefill.seq_lens
2306
2307 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
2314 - assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND"
2308 + assert get_flashinfer_layout_string(kv_cache_layout) == "HND"
2309 assert is_strictly_contiguous(prefill_query)
2310 assert is_strictly_contiguous(workspace_buffer)
2311 assert is_strictly_contiguous(block_tables_prefill)
@@ -2501,7 +2495,7 @@ class FlashInferImpl(AttentionImpl):
2495 # trtllm-gen needs HND layout on SM100. XQA is selected
2496 # separately on SM90 and does not use this SM100 layout gate.
2497 if decode_with_trtllm_gen:
2504 - assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND"
2498 + assert get_flashinfer_layout_string(kv_cache_layout) == "HND"
2499 else:
2500 assert decode_with_xqa
2501 assert is_strictly_contiguous(decode_query)
@@ -2547,7 +2541,7 @@ class FlashInferImpl(AttentionImpl):
2541 window_left=self.window_left,
2542 out=output_padded[:decode_query_tokens],
2543 sinks=self.sinks,
2550 - kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
2544 + kv_layout=get_flashinfer_layout_string(kv_cache_layout),
2545 q_len_per_req=q_len_per_req,
2546 mask=attn_metadata.decode.mask,
2547 q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens,
@@ -2614,7 +2608,7 @@ class FlashInferImpl(AttentionImpl):
2608 sinks=self.sinks,
2609 o_sf_scale=self.o_sf_scale,
2610 out=out,
2617 - kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
2611 + kv_layout=get_flashinfer_layout_string(kv_cache_layout),
2612 backend=attn_metadata.decode.kernel.value,
2613 q_len_per_req=q_len_per_req,
2614 max_q_len=max_q_len,
vllm/v1/attention/backends/utils.py
+22 -2
@@ -21,7 +21,12 @@ from vllm.config.cache import _layout_from_name
21 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
22 from vllm.utils.math_utils import cdiv
23 from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor
24 -from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MambaSpec
24 +from vllm.v1.kv_cache_interface import (
25 + AttentionSpec,
26 + KVCacheLayout,
27 + KVCacheSpec,
28 + MambaSpec,
29 +)
30
31 if TYPE_CHECKING:
32 from vllm.v1.core.sched.output import SchedulerOutput
@@ -286,9 +291,10 @@ def resolve_kv_cache_layout(
291 # specs can re-interpret HNC with different sizes as long as the total number of
292 # bytes is the same. If not block-compact, each spec must agree on HNC to alias
293 # the same page (this aliasing is done by the Hybrid Memory Allocator, HMA).
294 + kv_cache_specs = tuple(kv_cache_specs or ())
295 hnc_shapes = {
296 (spec.num_heads, spec.num_states, spec.page_size_bytes)
291 - for spec in kv_cache_specs or ()
297 + for spec in kv_cache_specs
298 }
299 if len(hnc_shapes) > 1:
300 candidates = [m for m in candidates if m.is_block_compact]
@@ -298,6 +304,20 @@ def resolve_kv_cache_layout(
304 f"none is in every supported set: {supported_layouts}."
305 )
306
307 + dcp_sharding = {
308 + spec.dcp_sharded for spec in kv_cache_specs if isinstance(spec, AttentionSpec)
309 + }
310 + page_sizes = {spec.page_size_bytes for spec in kv_cache_specs}
311 + if len(dcp_sharding) > 1 and len(page_sizes) > 1:
312 + # Sharded target and replicated draft caches need independent groups.
313 + # Block-outer layouts can pack those groups without equalizing pages.
314 + candidates = [layout for layout in candidates if layout.is_block_outermost]
315 + if not candidates:
316 + raise ValueError(
317 + "DCP with a replicated draft and mixed KV page sizes requires "
318 + f"a block-outer KV cache layout; supported sets: {supported_layouts}."
319 + )
320 +
321 if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None:
322 layout = _layout_from_name(requested)
323 if layout not in candidates:
vllm/v1/core/kv_cache_coordinator.py
+13 -13
@@ -11,7 +11,6 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
11 from vllm.v1.core.kv_cache_utils import (
12 BlockHash,
13 KVCacheBlock,
14 - dcp_world_size_for_kv_cache_spec,
14 )
15 from vllm.v1.core.single_type_kv_cache_manager import (
16 CrossAttentionManager,
@@ -20,6 +19,7 @@ from vllm.v1.core.single_type_kv_cache_manager import (
19 get_manager_for_kv_cache_spec,
20 )
21 from vllm.v1.kv_cache_interface import (
22 + AttentionSpec,
23 FullAttentionSpec,
24 KVCacheConfig,
25 KVCacheSpec,
@@ -140,9 +140,7 @@ class KVCacheCoordinator(ABC):
140 role=kv_cache_group.role,
141 enable_caching=enable_caching,
142 kv_cache_group_id=i,
143 - dcp_world_size=dcp_world_size_for_kv_cache_spec(
144 - kv_cache_group.kv_cache_spec, dcp_world_size
145 - ),
143 + dcp_world_size=dcp_world_size,
144 pcp_world_size=pcp_world_size,
145 scheduler_block_size=self.scheduler_block_size,
146 needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing,
@@ -664,16 +662,18 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
662 f"hash_block_size. block_sizes={cacheable_block_sizes}, "
663 f"hash_block_size={hash_block_size}"
664 )
667 - assert pcp_world_size == 1, "PCP not support hybrid attn now."
668 - if dcp_world_size > 1:
669 - # DCP shards full-attention KV across ranks and replicates Mamba
670 - # state; other spec types (e.g. sliding window) have no DCP-aware
671 - # handling yet, so reject them explicitly.
672 - for g in kv_cache_config.kv_cache_groups:
673 - assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), (
665 + for g in kv_cache_config.kv_cache_groups:
666 + spec = g.kv_cache_spec
667 + replicated = not spec.dcp_sharded
668 + if pcp_world_size > 1:
669 + assert isinstance(spec, FullAttentionSpec) or (
670 + isinstance(spec, AttentionSpec) and replicated
671 + ), "PCP only supports full attention and replicated draft groups."
672 + if dcp_world_size > 1:
673 + assert isinstance(spec, FullAttentionSpec) or replicated, (
674 "DCP with hybrid KV cache layouts only supports "
675 - "full-attention and Mamba groups, got: "
676 - f"{type(g.kv_cache_spec).__name__}."
675 + "full-attention, Mamba, and replicated draft groups, got: "
676 + f"{type(spec).__name__}."
677 )
678 # Fine-grained hash hits require Mamba "align" and compatible cache
679 # managers in every group. TP needs hashing finer than the Mamba block;
vllm/v1/core/kv_cache_utils.py
+19 -41
@@ -682,12 +682,7 @@ def hash_block_tokens(
682
683 def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
684 """Return the token span of a cache block under DCP."""
685 - layer_specs = iter_layer_specs(spec)
686 - if len(layer_specs) > 0 and all(
687 - isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs
688 - ):
689 - return spec.block_size * dcp_world_size
690 - return spec.block_size
685 + return spec.block_size * (dcp_world_size if spec.dcp_sharded else 1)
686
687
688 def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:
@@ -707,28 +702,6 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache
702 return replace(spec, block_size=block_size)
703
704
710 -def dcp_world_size_for_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> int:
711 - """Return the DCP size that owns this group's block geometry.
712 -
713 - Full-attention KV (including MLA) is sharded across DCP ranks, so prefix
714 - hashing and manager ``block_size`` use the process DCP size. Other specs
715 - keep replicated per-rank state (Mamba, sliding window, chunked-local) and
716 - must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
717 -
718 - Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
719 - ``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated
720 - draft group would need a different spec, not this helper.
721 - """
722 - if dcp_world_size <= 1:
723 - return 1
724 - inner = spec
725 - if isinstance(spec, UniformTypeKVCacheSpecs):
726 - inner = next(iter(spec.kv_cache_specs.values()))
727 - if isinstance(inner, FullAttentionSpec):
728 - return dcp_world_size
729 - return 1
730 -
731 -
705 def resolve_kv_cache_block_sizes(
706 kv_cache_config: KVCacheConfig,
707 vllm_config: VllmConfig,
@@ -752,6 +725,8 @@ def resolve_kv_cache_block_sizes(
725 groups = kv_cache_config.kv_cache_groups
726
727 if len(groups) <= 1:
728 + if groups and not groups[0].kv_cache_spec.dcp_sharded:
729 + dcp = 1
730 bs = cache_config.block_size * dcp
731 return bs, bs
732
@@ -1935,18 +1910,21 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1910 kv_cache_spec: The kv cache spec of each attention layer in the model
1911
1912 """
1938 - if is_kv_cache_spec_uniform(
1939 - kv_cache_spec
1940 - ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):
1941 - return
1942 -
1943 - logger.warning(
1944 - "Hybrid KV cache manager is disabled for this hybrid model, "
1945 - "This means we do not enable any optimizations for saving KV cache "
1946 - "memory (e.g., dropping the KV cache outside the sliding window). "
1947 - "The compute of layers like sliding window is still saved."
1948 - )
1949 - kv_cache_spec.update(_promote_local_kv_cache_specs(kv_cache_spec))
1913 + groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict)
1914 + for name, spec in kv_cache_spec.items():
1915 + replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded
1916 + groups[replicated][name] = spec
1917 + for specs in groups.values():
1918 + promoted_specs = _promote_local_kv_cache_specs(specs)
1919 + if promoted_specs == specs:
1920 + continue
1921 + logger.warning(
1922 + "Hybrid KV cache manager is disabled for this hybrid model, "
1923 + "This means we do not enable any optimizations for saving KV cache "
1924 + "memory (e.g., dropping the KV cache outside the sliding window). "
1925 + "The compute of layers like sliding window is still saved."
1926 + )
1927 + kv_cache_spec.update(promoted_specs)
1928
1929
1930 def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
@@ -2276,7 +2254,7 @@ def _ensure_min_page_size(
2254 scaled: list[KVCacheGroupSpec] = []
2255 for g in groups:
2256 s = g.kv_cache_spec
2279 - kw: dict[str, int] = {"block_size": s.block_size * scale}
2257 + kw: dict[str, Any] = {"block_size": s.block_size * scale}
2258 if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None:
2259 kw["page_size_padded"] = s.page_size_padded * scale
2260 scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw)))
vllm/v1/core/single_type_kv_cache_manager.py
+3 -9
@@ -106,11 +106,9 @@ class SingleTypeKVCacheManager(ABC):
106 # managers have been validated by the coordinator.
107 self.cache_hit_alignment_tokens = scheduler_block_size
108 # The block size for this manager; used for actual block allocation.
109 - self.block_size = kv_cache_spec.block_size
110 - self.dcp_world_size = dcp_world_size
111 - self.pcp_world_size = pcp_world_size
112 - if dcp_world_size > 1:
113 - self.block_size *= dcp_world_size
109 + self.dcp_world_size = dcp_world_size if kv_cache_spec.dcp_sharded else 1
110 + self.pcp_world_size = pcp_world_size if kv_cache_spec.dcp_sharded else 1
111 + self.block_size = kv_cache_spec.block_size * self.dcp_world_size
112 self.kv_cache_spec = kv_cache_spec
113 self.block_pool = block_pool
114 self.enable_caching = enable_caching
@@ -1455,10 +1453,6 @@ class MambaManager(SingleTypeKVCacheManager):
1453 self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs
1454 ) -> None:
1455 super().__init__(kv_cache_spec, block_pool, **kwargs)
1458 - # Mamba layers use TP instead of DCP, so each rank holds the full
1459 - # recurrent state. Undo the DCP/PCP block_size scaling that the base
1460 - # class applies for attention groups whose KV cache is partitioned.
1461 - self.block_size = kv_cache_spec.block_size
1456 self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
1457 self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks
1458 self.has_prefill_checkpoint_blocks = (
vllm/v1/kv_cache_interface.py
+24 -5
@@ -158,6 +158,8 @@ class KVCacheSpec:
158
159 # number of tokens in a block
160 block_size: int
161 + dcp_sharded: bool = field(default=False, kw_only=True)
162 + """Whether DCP shards this cache's token positions across ranks."""
163
164 block_stride_alignment: int | None = field(default=None, kw_only=True)
165 """Required byte alignment between physical blocks, including packed layers."""
@@ -481,6 +483,7 @@ class HiSparseResidentSpec(KVCacheSpec):
483
484 @dataclass(frozen=True, kw_only=True)
485 class AttentionSpec(KVCacheSpec):
486 + dcp_sharded: bool = True
487 num_kv_heads: int
488 head_size: int
489 dtype: torch.dtype
@@ -536,7 +539,9 @@ class AttentionSpec(KVCacheSpec):
539
540 def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
541 parallel_config = vllm_config.parallel_config
539 - kv_shard_count = parallel_config.decode_context_parallel_size
542 + kv_shard_count = (
543 + parallel_config.decode_context_parallel_size if self.dcp_sharded else 1
544 + )
545 return cdiv(max_len, self.block_size * kv_shard_count)
546
547
@@ -568,7 +573,7 @@ class FullAttentionSpec(AttentionSpec):
573 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
574 max_model_len = vllm_config.model_config.max_model_len
575 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
571 - if dcp_world_size > 1:
576 + if self.dcp_sharded and dcp_world_size > 1:
577 max_model_len = cdiv(max_model_len, dcp_world_size)
578 return cdiv(max_model_len, self.block_size) * self.page_size_bytes
579
@@ -612,6 +617,7 @@ class FullAttentionSpec(AttentionSpec):
617 head_size_v=specs[0].head_size_v,
618 dtype=specs[0].dtype,
619 kv_quant_mode=specs[0].kv_quant_mode,
620 + dcp_sharded=specs[0].dcp_sharded,
621 page_size_padded=specs[0].page_size_padded,
622 num_head_slots=specs[0].num_head_slots,
623 state_content_bytes=specs[0].state_content_bytes,
@@ -698,6 +704,7 @@ class MLAAttentionSpec(FullAttentionSpec):
704 head_size=specs[0].head_size,
705 dtype=specs[0].dtype,
706 kv_quant_mode=specs[0].kv_quant_mode,
707 + dcp_sharded=specs[0].dcp_sharded,
708 page_size_padded=specs[0].page_size_padded,
709 num_head_slots=specs[0].num_head_slots,
710 state_content_bytes=specs[0].state_content_bytes,
@@ -760,6 +767,7 @@ class RSWASpec(FullAttentionSpec):
767 head_size_v=base.head_size_v,
768 dtype=base.dtype,
769 kv_quant_mode=base.kv_quant_mode,
770 + dcp_sharded=base.dcp_sharded,
771 page_size_padded=base.page_size_padded,
772 num_head_slots=base.num_head_slots,
773 state_content_bytes=base.state_content_bytes,
@@ -850,9 +858,10 @@ class SlidingWindowSpec(AttentionSpec):
858 return cdiv(num_tokens, self.block_size) + 1
859
860 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
853 - assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
854 - "DCP not support sliding window."
855 - )
861 + assert (
862 + not self.dcp_sharded
863 + or vllm_config.parallel_config.decode_context_parallel_size == 1
864 + ), "DCP not support sliding window."
865 max_blocks = self.max_admission_blocks_per_request(
866 max_in_flight_tokens=vllm_config.max_in_flight_tokens,
867 max_model_len=vllm_config.model_config.max_model_len,
@@ -950,6 +959,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
959 extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
960 bounded_replay_set = set(spec.bounded_replay for spec in specs)
961 block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
962 + assert len({spec.dcp_sharded for spec in specs}) == 1
963 assert (
964 len(cache_dtype_str_set) == 1
965 and len(tokens_per_state_set) == 1
@@ -969,6 +979,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
979 num_kv_heads=specs[0].num_kv_heads,
980 head_size=specs[0].head_size,
981 dtype=specs[0].dtype,
982 + dcp_sharded=specs[0].dcp_sharded,
983 page_size_padded=specs[0].page_size_padded,
984 num_head_slots=specs[0].num_head_slots,
985 state_content_bytes=specs[0].state_content_bytes,
@@ -1185,6 +1196,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
1196 block_stride_alignment=specs[0].block_stride_alignment,
1197 dtype=specs[0].dtype,
1198 kv_quant_mode=specs[0].kv_quant_mode,
1199 + dcp_sharded=specs[0].dcp_sharded,
1200 page_size_padded=specs[0].page_size_padded,
1201 num_head_slots=specs[0].num_head_slots,
1202 state_content_bytes=specs[0].state_content_bytes,
@@ -1217,6 +1229,11 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1229
1230 kv_cache_specs: dict[str, KVCacheSpec]
1231
1232 + def __post_init__(self):
1233 + super().__post_init__()
1234 + if self.kv_cache_specs:
1235 + object.__setattr__(self, "dcp_sharded", self.first_spec.dcp_sharded)
1236 +
1237 @property
1238 def prefix_cacheable(self) -> bool:
1239 return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
@@ -1270,6 +1287,8 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1287 if len(block_sizes) > 1:
1288 # Different block sizes, not uniform.
1289 return False
1290 + if len({spec.dcp_sharded for spec in kv_cache_specs.values()}) > 1:
1291 + return False
1292 first_spec = next(iter(kv_cache_specs.values()))
1293 return first_spec.is_uniform_with_collection(kv_cache_specs)
1294
vllm/v1/simple_kv_offload/manager.py
+2 -6
@@ -27,13 +27,13 @@ from vllm.v1.core.kv_cache_coordinator import (
27 from vllm.v1.core.kv_cache_utils import (
28 BlockHashWithGroupId,
29 ExternalBlockHash,
30 - dcp_world_size_for_kv_cache_spec,
30 generate_block_hash_extra_keys,
31 get_block_hash,
32 get_group_id,
33 make_block_hash_with_group_id,
34 maybe_convert_block_hash,
35 resolve_block_hashes,
36 + resolve_dcp_kv_block_size,
37 )
38 from vllm.v1.core.sched.output import SchedulerOutput
39 from vllm.v1.kv_cache_interface import (
@@ -302,11 +302,7 @@ class SimpleCPUOffloadScheduler:
302 target = 0
303 for g in kv_cache_config.prefix_cacheable_groups:
304 spec = g.kv_cache_spec
305 - # Only full attention is sharded across DCP ranks; replicated specs
306 - # (mamba, sliding window, chunked-local) keep their own block size.
307 - block_size = spec.block_size * dcp_world_size_for_kv_cache_spec(
308 - spec, cp_world_size
309 - )
305 + block_size = resolve_dcp_kv_block_size(spec, cp_world_size)
306 if isinstance(spec, MambaSpec):
307 target += 2
308 elif isinstance(spec, SlidingWindowSpec):
vllm/v1/spec_decode/draft_model.py
+8 -1
@@ -81,13 +81,20 @@ class DraftModelProposer(SpecDecodeBaseProposer):
81 def _create_draft_vllm_config(self) -> VllmConfig:
82 base = super()._create_draft_vllm_config()
83 spec = self.speculative_config
84 + target_parallel = base.parallel_config
85
86 return replace(
87 base,
88 quant_config=None,
89 parallel_config=replace(
90 spec.draft_parallel_config,
90 - rank=self.vllm_config.parallel_config.rank,
91 + rank=target_parallel.rank,
92 + decode_context_parallel_size=(
93 + target_parallel.decode_context_parallel_size
94 + ),
95 + cp_kv_cache_interleave_size=target_parallel.cp_kv_cache_interleave_size,
96 + dcp_comm_backend=target_parallel.dcp_comm_backend,
97 + dcp_q_replicate=target_parallel.dcp_q_replicate,
98 ),
99 model_config=spec.draft_model_config,
100 )
vllm/v1/worker/cp_utils.py
+2
@@ -40,6 +40,8 @@ def check_attention_cp_compatibility(
40 layer_impl = getattr(layer, "impl", None)
41 if layer_impl is None:
42 continue
43 + if not check_pcp and layer_impl.dcp_world_size == 1:
44 + continue
45 if vllm_config.speculative_config is not None and interleave_size > 1:
46 assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, (
47 "MTP with cp_kv_cache_interleave_size > 1 is not "
vllm/v1/worker/gpu/block_table.py
+11
@@ -27,6 +27,7 @@ class BlockTables:
27 cp_rank: int = 0,
28 cp_interleave: int = 1,
29 slot_mapping_enabled: list[bool] | None = None,
30 + dcp_sharded: list[bool] | None = None,
31 ):
32 self.block_sizes = block_sizes
33 self.kernel_block_sizes = kernel_block_sizes
@@ -44,6 +45,10 @@ class BlockTables:
45 slot_mapping_enabled = [True] * self.num_kv_cache_groups
46 assert len(slot_mapping_enabled) == self.num_kv_cache_groups
47 self._slot_mapping_enabled = slot_mapping_enabled
48 + if dcp_sharded is None:
49 + dcp_sharded = [True] * self.num_kv_cache_groups
50 + assert len(dcp_sharded) == self.num_kv_cache_groups
51 + self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device)
52
53 self.blocks_per_kv_block = [
54 bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
@@ -210,6 +215,7 @@ class BlockTables:
215 self.block_sizes_tensor,
216 self.kernel_block_sizes_tensor,
217 self.slot_mapping_enabled,
218 + self.dcp_sharded,
219 slot_mappings,
220 slot_mappings.stride(0),
221 self.cp_rank,
@@ -283,6 +289,7 @@ def _compute_slot_mappings_kernel(
289 block_sizes, # [num_kv_cache_groups]
290 kernel_block_sizes, # [num_kv_cache_groups]
291 slot_mapping_enabled, # [num_kv_cache_groups]
292 + dcp_sharded, # [num_kv_cache_groups]
293 slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
294 slot_mappings_stride,
295 cp_rank,
@@ -312,6 +319,8 @@ def _compute_slot_mappings_kernel(
319 kv_block_size = tl.load(block_sizes + group_id)
320 kernel_block_size = tl.load(kernel_block_sizes + group_id)
321 mapping_enabled = tl.load(slot_mapping_enabled + group_id)
322 + if CP_SIZE != 1:
323 + sharded = tl.load(dcp_sharded + group_id)
324
325 req_state_idx = tl.load(idx_mapping + batch_idx)
326 # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns
@@ -337,6 +346,8 @@ def _compute_slot_mappings_kernel(
346 remainder = virtual_block_offsets % CP_INTERLEAVE
347 local_offsets = rounds * CP_INTERLEAVE + remainder
348 local_positions = virtual_block_indices * kv_block_size + local_offsets
349 + local_positions = tl.where(sharded, local_positions, positions)
350 + is_local = ~sharded | is_local
351
352 block_indices = tl.where(
353 mapping_enabled, local_positions // kernel_block_size, 0
vllm/v1/worker/gpu/model_runner.py
+14 -1
@@ -21,6 +21,7 @@ import gc
21 import time
22 from contextlib import AbstractContextManager
23 from copy import deepcopy
24 +from dataclasses import replace
25 from typing import Any, NamedTuple
26
27 import numpy as np
@@ -561,7 +562,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
562 return encoder_runner.get_encoder_timing_stats()
563
564 def get_kv_cache_spec(self):
564 - return get_kv_cache_spec(self.vllm_config)
565 + kv_cache_spec = get_kv_cache_spec(self.vllm_config)
566 + # A draft running at DCP=1 under a DCP target keeps a replicated cache.
567 + if (
568 + isinstance(self.speculator, DraftModelSpeculator)
569 + and self.speculator.dcp_size == 1
570 + and self.dcp_size > 1
571 + ):
572 + for name in self.speculator.draft_attn_layer_names & kv_cache_spec.keys():
573 + kv_cache_spec[name] = replace(kv_cache_spec[name], dcp_sharded=False)
574 + return kv_cache_spec
575
576 def initialize_kv_cache(
577 self,
@@ -587,6 +597,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
597 block_sizes = []
598 max_num_blocks_per_group = []
599 slot_mapping_enabled = []
600 + dcp_sharded = []
601 for kv_cache_group in kv_cache_config.kv_cache_groups:
602 spec = kv_cache_group.kv_cache_spec
603 block_sizes.append(spec.block_size)
@@ -594,6 +605,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
605 spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
606 )
607 slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
608 + dcp_sharded.append(spec.dcp_sharded)
609 # Let each cache type account for CP. Attention KV is DCP-sharded,
610 # while Mamba/GDN recurrent state is replicated across DCP ranks.
611 max_num_blocks = spec.max_num_blocks_per_req(
@@ -657,6 +669,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
669 device=self.device,
670 kernel_block_sizes=self.kernel_block_sizes,
671 slot_mapping_enabled=slot_mapping_enabled,
672 + dcp_sharded=dcp_sharded,
673 cp_size=self.dcp_size,
674 cp_rank=self.dcp_rank,
675 cp_interleave=self.cp_interleave,
vllm/v1/worker/gpu/pcp_manager.py
+4 -1
@@ -145,7 +145,10 @@ class PCPManager:
145 if speculative_config is not None:
146 if speculative_config.use_dspark():
147 dcp_size = parallel_config.decode_context_parallel_size
148 - if dcp_size not in (1, pcp_size):
148 + if (
149 + dcp_size not in (1, pcp_size)
150 + and speculative_config.draft_model_config.use_mla
151 + ):
152 raise NotImplementedError(
153 "MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got "
154 f"DCP={dcp_size}, PCP={pcp_size}."
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+13 -7
@@ -37,12 +37,18 @@ class DFlashSpeculator(DraftModelSpeculator):
37
38 def __init__(self, vllm_config: VllmConfig, device: torch.device):
39 parallel_config = vllm_config.parallel_config
40 - if parallel_config.prefill_context_parallel_size > 1:
41 - vllm_config = copy.copy(vllm_config)
42 - vllm_config.parallel_config = replace(
43 - parallel_config,
44 - prefill_context_parallel_size=1,
45 - )
40 + speculative_config = vllm_config.speculative_config
41 + assert speculative_config is not None
42 + vllm_config = copy.copy(vllm_config)
43 + vllm_config.parallel_config = replace(
44 + parallel_config,
45 + prefill_context_parallel_size=1,
46 + decode_context_parallel_size=(
47 + parallel_config.decode_context_parallel_size
48 + if speculative_config.draft_model_config.use_mla
49 + else 1
50 + ),
51 + )
52 super().__init__(vllm_config, device)
53
54 self.hidden_states = torch.zeros(
@@ -427,7 +433,7 @@ class DFlashSpeculator(DraftModelSpeculator):
433 self.block_tables.input_block_tables[gid],
434 self.block_tables.kernel_block_sizes[gid],
435 self.block_tables.cp_rank,
430 - self.block_tables.cp_size,
436 + self.dcp_size,
437 self.block_tables.cp_interleave,
438 self.parallel_drafting_token_id,
439 self.num_query_per_req,
vllm/v1/worker/gpu/spec_decode/speculator.py
+3 -2
@@ -138,6 +138,7 @@ class DraftModelSpeculator(BaseSpeculator):
138 # DP configuration
139 self.dp_size = vllm_config.parallel_config.data_parallel_size
140 self.dp_rank = vllm_config.parallel_config.data_parallel_rank
141 + self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
142
143 self.eplb_state: EplbState | None = None
144
@@ -332,14 +333,14 @@ class DraftModelSpeculator(BaseSpeculator):
333 out=draft_seq_lens_cpu_upper_bound[:num_reqs],
334 )
335 draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
335 - if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
336 + if dcp_local_seq_lens is None and self.dcp_size > 1:
337 # Draft steps advance and rewind their own global sequence lengths,
338 # so the target model's DCP-local lengths may already be stale.
339 dcp_local_seq_lens = prepare_dcp_local_seq_lens(
340 self.input_buffers.dcp_local_seq_lens,
341 self.input_buffers.seq_lens,
342 num_reqs,
342 - self.block_tables.cp_size,
343 + self.dcp_size,
344 self.block_tables.cp_rank,
345 self.block_tables.cp_interleave,
346 )