iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Lucas Wilkinson
2026-09-26 18:12:35 -04:00
committed by GitHub
co-authored by Claude Opus 5 Codex
parent 77871126f9
commit 927c87b347
32 changed files with 606 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_attention_backends.py
+1 -2
View File
@@ -942,8 +942,6 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa
impl.o_sf_scale = None
impl.window_left = -1
impl.sinks = None
impl.cache_config = unittest.mock.Mock()
impl.cache_config.get_resolved_kv_cache_layout.return_value = KVCacheLayout.LBHNC
layer = unittest.mock.Mock(
_q_scale=torch.tensor(1.0),
@@ -959,6 +957,7 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa
q_len_per_req=1,
)
attn_metadata = flashinfer_backend.FlashInferMetadata(
kv_cache_layout=KVCacheLayout.LBHNC,
num_actual_tokens=2,
slot_mapping=torch.empty(0, dtype=torch.int64),
q_data_type_prefill=torch.bfloat16,
tests/v1/attention/test_group_head_counts.py
+1
View File
@@ -101,6 +101,7 @@ def test_flash_attention_geometry_comes_from_the_group():
num_kv_heads=2,
head_size=64,
dtype=torch.bfloat16,
dcp_sharded=True,
)
with (
tests/v1/core/test_kv_cache_utils.py
+143 -12
View File
@@ -1599,19 +1599,142 @@ 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
@pytest.mark.parametrize("sliding_window", [None, 256])
@pytest.mark.parametrize("disable_hybrid", [False, True])
@pytest.mark.parametrize("pcp_size", [1, 4])
def test_dcp_target_allocates_replicated_draft_independently(
monkeypatch, sliding_window, disable_hybrid, pcp_size
):
"""A draft must retain all positions even when the target shards them."""
from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
config.parallel_config.decode_context_parallel_size = 4
config.parallel_config.prefill_context_parallel_size = pcp_size
config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hybrid
config.cache_config.block_size = 16
config.cache_config.kv_cache_layout = None
draft_args = dict(
block_size=16,
num_kv_heads=8,
head_size=64,
dtype=torch.bfloat16,
dcp_sharded=False,
)
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
draft = (
FullAttentionSpec(**draft_args)
if sliding_window is None
else SlidingWindowSpec(**draft_args, sliding_window=sliding_window)
)
specs = {"target": new_mla_spec(), "draft": draft}
layout = resolve_kv_cache_layout(config, [["LBHNC", "BLHNC"]], specs.values())
assert layout == KVCacheLayout.BLHNC
groups = get_kv_cache_groups(config, specs)
assert [group.layer_names for group in groups] == [["target"], ["draft"]]
widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups]
assert widths == [16, 64]
cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
config, groups, available_memory=16 * 1024 * 1024
)
scheduler_config = generate_scheduler_kv_cache_config([cache_config])
assert get_max_concurrency_for_kv_cache_config(config, cache_config) > 0
assert get_max_concurrency_for_kv_cache_config(config, cache_config) == (
get_max_concurrency_for_kv_cache_config(config, scheduler_config)
)
manager = KVCacheManager(
scheduler_config,
max_model_len=1024,
hash_block_size=16,
scheduler_block_size=64,
dcp_world_size=4,
pcp_world_size=pcp_size,
enable_caching=True,
)
assert manager.coordinator.group_block_sizes == (64, 16)
request = make_request("replicated-draft", [1] * 65, block_size=16, hash_fn=sha256)
blocks = manager.allocate_slots(request, 65)
assert blocks is not None
assert [len(group) for group in blocks.blocks] == [2, 5]
manager.cache_blocks(request, 64)
cached_request = make_request(
"cached-draft", [1] * 65, block_size=16, hash_fn=sha256
)
cached_blocks, num_cached, _ = manager.get_computed_blocks(cached_request)
assert num_cached == 64
assert [len(group) for group in cached_blocks.blocks] == [1, 4]
@pytest.mark.parametrize("use_mla", [False, True])
def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla):
config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
config.parallel_config.decode_context_parallel_size = 4
draft = replace(
new_mla_spec() if use_mla else new_kv_cache_spec(),
block_size=16,
dcp_sharded=False,
)
merged = type(draft).merge([draft, draft])
wrapped = UniformTypeKVCacheSpecs.from_specs({"draft": merged})
assert wrapped is not None and not wrapped.dcp_sharded
assert merged.max_num_blocks_per_req(config, 1024) == 64
assert merged.max_memory_usage_bytes(config) == 64 * draft.page_size_bytes
assert kv_cache_utils.resolve_dcp_kv_block_size(merged, 4) == 16
with pytest.raises(AssertionError):
type(draft).merge([draft, replace(draft, dcp_sharded=True)])
@pytest.mark.parametrize("with_draft", [False, True])
@pytest.mark.parametrize("indexer_alignment", [None, 512])
def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment):
from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view
from vllm.v1.worker.utils import allocate_kv_cache
config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
config.cache_config.kv_cache_layout = "BLHNC"
config.parallel_config.decode_context_parallel_size = 4
common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8)
specs = {
"target": MLAAttentionSpec(
**common,
head_size=576,
state_content_bytes=656,
block_stride_alignment=656,
cache_dtype_str="fp8_ds_mla",
is_index_group_leader=True,
),
"indexer": MLAAttentionSpec(
**common,
head_size=132,
cache_role=SparseCacheRole.INDEXER,
block_stride_alignment=indexer_alignment,
),
"draft": SlidingWindowSpec(
block_size=64,
num_kv_heads=64,
head_size=64,
dtype=torch.bfloat16,
sliding_window=2048,
dcp_sharded=False,
),
}
if not with_draft:
del specs["draft"]
groups = get_kv_cache_groups(config, specs)
cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
config, groups, available_memory=8 * 1024 * 1024
)
caches = allocate_kv_cache(
cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups)
)
cache = caches["target"].squeeze(1)
if indexer_alignment is not None:
assert caches["indexer"].stride(0) % indexer_alignment == 0
rows, stride_rows = flat_kv_row_view(cache, 64)
cache[1, 0].fill_(7)
torch.testing.assert_close(rows[stride_rows], cache[1, 0])
assert torch.all(rows[stride_rows] == 7)
@pytest.mark.parametrize(
@@ -3815,6 +3938,8 @@ def test_unify_hybrid_kv_cache_specs():
"layer_1": before_spec_1,
"layer_2": before_spec_2,
}
kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False)
kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False)
kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
expected_spec_1 = new_kv_cache_spec(block_size=64)
expected_spec_2 = new_kv_cache_spec(
@@ -3823,6 +3948,7 @@ def test_unify_hybrid_kv_cache_specs():
assert kv_cache_spec["layer_1"] == expected_spec_1
assert kv_cache_spec["layer_2"] == expected_spec_2
assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024
assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False)
# 2. has_full_attention and has_chunked_local_attention
before_spec_1 = new_kv_cache_spec()
@@ -3874,6 +4000,11 @@ def test_unify_hybrid_kv_cache_specs():
with pytest.raises(ValueError):
kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
# Replicated Mamba state still requires the hybrid cache manager.
kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()}
with pytest.raises(ValueError):
kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
def test_unify_kv_cache_spec_page_size_mamba():
"""Regression test for https://github.com/vllm-project/vllm/issues/43626.
tests/v1/core/test_prefix_caching.py
+12 -7
View File
@@ -1336,12 +1336,12 @@ def make_kv_cache_config_three_types(
)
def test_prefix_cache_hit_uses_per_group_dcp_geometry():
@pytest.mark.parametrize("draft_sharded", [True, False])
def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded):
"""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.
Target MLA is sharded; draft MLA can opt out and Mamba stays replicated.
Hits must align to the target block in either case.
"""
block_size = 16
dcp = 8
@@ -1362,6 +1362,7 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry():
KVCacheGroupSpec(
["draft_mla"],
MLAAttentionSpec(
dcp_sharded=draft_sharded,
block_size=block_size,
num_kv_heads=1,
head_size=1,
@@ -1388,10 +1389,10 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry():
)
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 draft_mgr.dcp_world_size == (dcp if draft_sharded else 1)
assert mamba_mgr.dcp_world_size == 1
assert target_mgr.block_size == sharded_block
assert draft_mgr.block_size == sharded_block
assert draft_mgr.block_size == (sharded_block if draft_sharded else block_size)
assert mamba_mgr.block_size == block_size
hash_fn = sha256
@@ -1407,7 +1408,11 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry():
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]
assert [len(group) for group in computed_blocks.blocks] == [
2,
2 if draft_sharded else 16,
16,
]
manager.free(req0)
manager.free(req1)
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
+33
View File
@@ -11,6 +11,7 @@ from vllm.config import CompilationConfig, VllmConfig, replace
from vllm.config.kernel import MoEBackend
from vllm.engine.arg_utils import EngineArgs
from vllm.platforms import current_platform
from vllm.v1.spec_decode.draft_model import DraftModelProposer
from ..utils import (
Messages,
@@ -196,6 +197,38 @@ def test_draft_model_engine_args_tensor_parallelism():
assert draft_config.quant_config is None
@pytest.mark.cpu_test
@pytest.mark.parametrize("dcp_size", [1, 2])
def test_draft_model_inherits_target_dcp_config(monkeypatch, dcp_size):
monkeypatch.setattr(current_platform, "device_count", lambda: 16)
target_config = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=16,
decode_context_parallel_size=dcp_size,
cp_kv_cache_interleave_size=16,
dcp_comm_backend="ag_rs",
distributed_executor_backend="mp",
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
"num_speculative_tokens": 3,
},
).create_engine_config()
proposer = object.__new__(DraftModelProposer)
proposer.vllm_config = target_config
proposer.speculative_config = target_config.speculative_config
draft_config = proposer._create_draft_vllm_config()
draft_parallel = draft_config.parallel_config
assert draft_parallel.decode_context_parallel_size == dcp_size
assert draft_parallel.cp_kv_cache_interleave_size == 16
assert draft_parallel.dcp_comm_backend == "ag_rs"
assert (
draft_parallel.dcp_q_replicate == target_config.parallel_config.dcp_q_replicate
)
assert draft_parallel is not target_config.parallel_config
assert target_config.parallel_config.decode_context_parallel_size == dcp_size
def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
"""Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic
so we can test it without instantiating a full proposer."""
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+2 -1
View File
@@ -500,7 +500,7 @@ def test_decode_bench_connector_skips_null_padding():
def test_decode_bench_connector_uses_per_group_dcp_block_sizes():
"""DCP scales full-attention blocks but not sliding-window blocks."""
"""DCP scales target blocks but not the replicated draft's window."""
block_size = 16
dcp_world_size = 2
vllm_config = create_vllm_config(
@@ -515,6 +515,7 @@ def test_decode_bench_connector_uses_per_group_dcp_block_sizes():
dtype=torch.float32,
)
sliding_window = SlidingWindowSpec(
dcp_sharded=False,
block_size=block_size,
num_kv_heads=1,
head_size=1,
tests/v1/simple_kv_offload/test_scheduler.py
+2 -2
View File
@@ -2316,8 +2316,8 @@ def _make_hybrid_attention_mamba_scheduler(
def test_hybrid_store_uses_resolved_group_block_sizes() -> None:
"""Replicated groups must not be scaled by the DCP world size.
``dcp_world_size_for_kv_cache_spec`` gives full attention the process DCP
size and every other spec 1. Reconstructing a group's geometry as
The spec's ``dcp_sharded`` flag controls its effective block size.
Reconstructing a group's geometry as
``spec.block_size * cp_world_size`` over-scales the replicated groups, so
the eager store scan believes far fewer of their blocks are ready and
silently offloads only a fraction of them.
tests/v1/spec_decode/test_dflash_prepare_inputs.py
+68 -30
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import pytest
@@ -9,7 +10,7 @@ import torch
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
prepare_dflash_inputs,
DFlashSpeculator,
)
pytestmark = pytest.mark.skipif(
@@ -24,6 +25,7 @@ def _run_prepare(
cp_rank: int = 0,
cp_size: int = 1,
cp_interleave: int = 1,
draft_dcp_size: int | None = None,
):
device = torch.device("cuda")
max_num_reqs = 4
@@ -40,7 +42,9 @@ def _run_prepare(
)
input_batch = SimpleNamespace(
num_reqs=1,
num_tokens=4,
num_scheduled_tokens=np.array([4], dtype=np.int32),
seq_lens_cpu_upper_bound=torch.tensor([target_positions[-1] + 1]),
positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),
@@ -74,36 +78,65 @@ def _run_prepare(
next_prefill_tokens = torch.zeros_like(last_sampled)
block_table = torch.tensor([block_table_values], dtype=torch.int32, device=device)
prepare_dflash_inputs(
input_buffers,
query_slot_mapping,
context_positions,
context_slot_mapping,
sample_indices,
sample_pos,
sample_idx_mapping,
temperature,
seeds,
input_batch,
torch.tensor([1], dtype=torch.int32, device=device),
torch.tensor([2], dtype=torch.int32, device=device),
last_sampled,
next_prefill_tokens,
input_temperature,
input_seeds,
block_table,
4,
cp_rank,
cp_size,
cp_interleave,
123,
num_speculative_steps,
num_speculative_steps,
max_num_reqs,
max_num_tokens,
128,
class InputsPrepared(Exception):
pass
draft = SimpleNamespace(
input_buffers=input_buffers,
context_positions=context_positions,
_context_slot_mappings=context_slot_mapping.unsqueeze(0),
sample_indices=sample_indices,
sample_pos=sample_pos,
sample_idx_mapping=sample_idx_mapping,
temperature=temperature,
seeds=seeds,
hidden_states=torch.zeros(4, 1, device=device),
prepare_context_anchor=lambda *args: None,
query_cudagraph_manager=None,
dp_size=1,
dp_rank=0,
pcp_manager=None,
draft_kv_cache_group_id=0,
draft_kv_cache_group_ids=[0],
_layer_group_idx=None,
block_tables=SimpleNamespace(
slot_mappings=query_slot_mapping.unsqueeze(0),
input_block_tables=[block_table],
kernel_block_sizes=[4],
cp_rank=cp_rank,
cp_size=cp_size,
cp_interleave=cp_interleave,
),
dcp_size=cp_size if draft_dcp_size is None else draft_dcp_size,
parallel_drafting_token_id=123,
num_query_per_req=num_speculative_steps,
num_speculative_steps=num_speculative_steps,
max_num_reqs=max_num_reqs,
max_num_tokens=max_num_tokens,
max_model_len=128,
sample_from_anchor=True,
)
with (
patch(
"vllm.v1.worker.gpu.spec_decode.dflash.speculator.dispatch_cg_and_sync_dp",
side_effect=InputsPrepared,
),
pytest.raises(InputsPrepared),
):
DFlashSpeculator.propose(
draft,
input_batch,
{},
{},
draft.hidden_states,
None,
torch.tensor([1], dtype=torch.int32, device=device),
torch.tensor([2], dtype=torch.int32, device=device),
last_sampled,
next_prefill_tokens,
input_temperature,
input_seeds,
)
torch.accelerator.synchronize()
return SimpleNamespace(
input_buffers=input_buffers,
@@ -118,12 +151,17 @@ def _run_prepare(
)
def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
@pytest.mark.parametrize("cp_rank,cp_size", [(0, 1), (0, 2), (1, 2)])
def test_prepare_dflash_inputs_excludes_rejected_context_suffix(cp_rank, cp_size):
# Positions 10/11 use physical block 7. Rejected positions 12/13 would use
# block 8, but must be PAD context rather than contaminating draft KV.
out = _run_prepare(
target_positions=[10, 11, 12, 13],
block_table_values=[0, 0, 7, 8, 9, 10, 11, 12],
cp_rank=cp_rank,
cp_size=cp_size,
cp_interleave=2,
draft_dcp_size=1,
)
assert out.context_positions[:4].tolist() == [10, 11, 0, 0]
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+9 -1
View File
@@ -17,6 +17,7 @@ from types import MethodType, SimpleNamespace
from unittest.mock import patch
import numpy as np
import pytest
import torch
from vllm.config.compilation import CUDAGraphMode
@@ -53,6 +54,7 @@ def _make_fake_speculator(
arange_np=np.arange(max_num_reqs + 1, dtype=np.int32),
draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool),
block_tables=fake_block_tables,
dcp_size=1,
input_buffers=fake_input_buffers,
attn_groups=[],
kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
@@ -152,14 +154,17 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len():
assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
@pytest.mark.parametrize("draft_dcp_size", [1, 2])
def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens(draft_dcp_size):
fake = _make_fake_speculator()
fake.dcp_size = draft_dcp_size
fake.block_tables.cp_size = 2
fake.block_tables.cp_rank = 1
fake.block_tables.cp_interleave = 4
fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16])
def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave):
assert draft_dcp_size > 1
assert seq_lens is fake.input_buffers.seq_lens
assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4)
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():
)
local = captured["dcp_local_seq_lens"]
if draft_dcp_size == 1:
assert local is None
return
assert isinstance(local, torch.Tensor)
assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr()
assert local.tolist() == [1, 4, 8, 0]
tests/v1/worker/test_cp_utils.py
+34
View File
@@ -1,5 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
@@ -7,6 +9,38 @@ from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention
def test_cp_checks_allow_a_replicated_draft_backend_without_dcp(monkeypatch):
from vllm.v1.worker import cp_utils
config = SimpleNamespace(
parallel_config=SimpleNamespace(
prefill_context_parallel_size=1,
decode_context_parallel_size=4,
cp_kv_cache_interleave_size=16,
),
speculative_config=object(),
)
target = SimpleNamespace(
impl=SimpleNamespace(
dcp_world_size=4,
supports_mtp_with_cp_non_trivial_interleave_size=True,
need_to_return_lse_for_decode=True,
)
)
draft = SimpleNamespace(
impl=SimpleNamespace(
dcp_world_size=1,
supports_mtp_with_cp_non_trivial_interleave_size=False,
need_to_return_lse_for_decode=False,
)
)
layers = {"target": target, "draft": draft}
monkeypatch.setattr(cp_utils, "get_layers_from_vllm_config", lambda *args: layers)
cp_utils.check_attention_cp_compatibility(config, target_layer_names={"target"})
with pytest.raises(AssertionError, match="MTP with cp_kv_cache_interleave_size"):
cp_utils.check_attention_cp_compatibility(config)
def test_skip_gate_only_for_zero_context():
assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32))
assert not should_skip_dcp_context_attention(
tests/v1/worker/test_dspark_utils.py
+92
View File
@@ -2,8 +2,17 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass, field
from types import SimpleNamespace
import pytest
import torch
from vllm.config import ParallelConfig
from vllm.v1.kv_cache_interface import FullAttentionSpec
from vllm.v1.worker.gpu import model_runner
from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config
from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
@dataclass
@@ -49,3 +58,86 @@ def test_dspark_parallel_config_disables_eplb_atomically():
assert draft_config.eplb_config.num_redundant_experts == 0
assert not draft_config.enable_elastic_ep
assert draft_config.eplb_config is not target_config.eplb_config
@pytest.mark.parametrize("pcp_size", [1, 4])
@pytest.mark.parametrize("dcp_size", [1, 4])
@pytest.mark.parametrize("use_mla", [False, True])
def test_draft_context_parallelism_without_changing_target(
monkeypatch, pcp_size, dcp_size, use_mla
):
target_parallel = ParallelConfig(
tensor_parallel_size=4,
prefill_context_parallel_size=pcp_size,
decode_context_parallel_size=dcp_size,
cp_kv_cache_interleave_size=16,
distributed_executor_backend="mp",
)
target_config = SimpleNamespace(
parallel_config=target_parallel,
speculative_config=SimpleNamespace(
draft_model_config=SimpleNamespace(use_mla=use_mla)
),
)
class CapturedConfig(Exception):
pass
def capture_init(self, config, device):
raise CapturedConfig(config)
monkeypatch.setattr(DraftModelSpeculator, "__init__", capture_init)
with pytest.raises(CapturedConfig) as captured:
DFlashSpeculator(target_config, device=None)
draft_parallel = captured.value.args[0].parallel_config
assert draft_parallel.tensor_parallel_size == 4
assert draft_parallel.prefill_context_parallel_size == 1
assert draft_parallel.decode_context_parallel_size == (dcp_size if use_mla else 1)
assert draft_parallel.cp_kv_cache_interleave_size == 16
assert target_parallel.prefill_context_parallel_size == pcp_size
assert target_parallel.decode_context_parallel_size == dcp_size
assert target_parallel.cp_kv_cache_interleave_size == 16
@pytest.mark.parametrize("configured_dcp", [1, 4])
def test_attention_uses_draft_dcp_setting_inside_target_process_group(
monkeypatch, configured_dcp
):
import vllm.config as config_module
from vllm.distributed import parallel_state
from vllm.v1.attention.backend import AttentionImplBase
config = SimpleNamespace(
parallel_config=SimpleNamespace(decode_context_parallel_size=configured_dcp)
)
monkeypatch.setattr(
config_module, "get_current_vllm_config_or_none", lambda: config
)
monkeypatch.setattr(
parallel_state, "_DCP", SimpleNamespace(world_size=4, rank_in_group=2)
)
impl = AttentionImplBase()
assert impl.dcp_world_size == configured_dcp
assert impl.dcp_rank == (0 if configured_dcp == 1 else 2)
@pytest.mark.parametrize("draft_dcp_size", [1, 4])
def test_runner_marks_only_replicated_draft_caches(monkeypatch, draft_dcp_size):
spec = FullAttentionSpec(
block_size=16, num_kv_heads=1, head_size=64, dtype=torch.bfloat16
)
monkeypatch.setattr(
model_runner, "get_kv_cache_spec", lambda _: {"target": spec, "draft": spec}
)
runner = object.__new__(model_runner.GPUModelRunner)
runner.vllm_config = None
runner.dcp_size = 4
runner.speculator = object.__new__(DFlashSpeculator)
runner.speculator.dcp_size = draft_dcp_size
runner.speculator.draft_attn_layer_names = {"draft"}
specs = runner.get_kv_cache_spec()
assert specs["target"].dcp_sharded
assert specs["draft"].dcp_sharded == (draft_dcp_size == 4)
tests/v1/worker/test_gpu_autoregressive_speculator.py
+2
View File
@@ -88,6 +88,7 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):
query_start_loc=torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32),
seq_lens=torch.tensor([11, 21, 0, 0], dtype=torch.int32),
)
speculator.dcp_size = 1
speculator.block_tables = SimpleNamespace(
cp_size=1,
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)
parallel_config=SimpleNamespace(
data_parallel_size=1,
data_parallel_rank=0,
decode_context_parallel_size=1,
),
)
tests/v1/worker/test_gpu_block_table.py
+9 -7
View File
@@ -171,22 +171,23 @@ def test_block_tables_skip_custom_slot_mapping_groups():
@pytest.mark.parametrize("cp_rank", range(4))
def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
"""DCP interleave is expressed in logical-block token coordinates."""
"""Only sharded groups use DCP interleave in logical-block coordinates."""
device = torch.device("cuda")
block_tables = BlockTables(
block_sizes=[128],
block_sizes=[128, 128],
max_num_reqs=1,
max_num_batched_tokens=1024,
max_num_blocks_per_group=[2],
max_num_blocks_per_group=[2, 8],
device=device,
kernel_block_sizes=[64],
kernel_block_sizes=[64, 64],
dcp_sharded=[True, False],
cp_size=4,
cp_rank=cp_rank,
cp_interleave=128,
)
block_tables.append_block_ids(
req_index=0,
new_block_ids=([5, 9],),
new_block_ids=([5, 9], list(range(10, 18))),
overwrite=True,
)
block_tables.apply_staged_writes()
@@ -199,7 +200,7 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
query_start_loc,
positions,
num_tokens_padded=1024,
)[0]
)
expected = torch.full((1024,), -1, dtype=torch.int64, device=device)
first_start = cp_rank * 128
@@ -210,7 +211,8 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
expected[second_start : second_start + 128] = torch.arange(
9 * 128, 10 * 128, dtype=torch.int64, device=device
)
assert torch.equal(actual, expected)
assert torch.equal(actual[0], expected)
assert torch.equal(actual[1], positions + 10 * 128)
def test_v1_block_table_move_row_clears_vacated_row():
tests/v1/worker/test_gpu_pcp_manager.py
+18
View File
@@ -75,6 +75,24 @@ def test_validate_config_rejects_full_graph_for_prefills():
)
@pytest.mark.parametrize("draft_uses_mla", [False, True])
@pytest.mark.parametrize("dcp_size", [1, 2, 4])
def test_validate_config_dcp_for_dspark(draft_uses_mla, dcp_size):
config = _make_config(CUDAGraphMode.NONE)
config.parallel_config.prefill_context_parallel_size = 4
config.parallel_config.decode_context_parallel_size = dcp_size
config.model_config.hf_text_config.index_topk = 1
config.speculative_config = SimpleNamespace(
use_dspark=lambda: True,
draft_model_config=SimpleNamespace(use_mla=draft_uses_mla),
)
if draft_uses_mla and dcp_size == 2:
with pytest.raises(NotImplementedError, match="requires DCP=1 or DCP=PCP"):
PCPManager.validate_config(config, supports_mm_inputs=False)
return
PCPManager.validate_config(config, supports_mm_inputs=False)
def test_replicated_decode_piecewise_graph_padding(monkeypatch):
manager = PCPManager(
pcp_world_size=2,
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+1 -8
View File
@@ -59,7 +59,6 @@ from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv
from vllm.v1.attention.backend import AttentionMetadata
from vllm.v1.core.kv_cache_utils import (
dcp_world_size_for_kv_cache_spec,
resolve_dcp_kv_block_size,
)
from vllm.v1.kv_cache_interface import (
@@ -259,13 +258,7 @@ class DecodeBenchConnectorScheduler:
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
self.kv_cache_groups = kv_cache_config.kv_cache_groups
self.group_block_sizes = tuple(
resolve_dcp_kv_block_size(
group.kv_cache_spec,
dcp_world_size_for_kv_cache_spec(
group.kv_cache_spec,
dcp_world_size,
),
)
resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_world_size)
for group in self.kv_cache_groups
)
vllm/distributed/parallel_state.py
+8
View File
@@ -1583,6 +1583,14 @@ def get_dcp_group() -> GroupCoordinator:
return _DCP
def get_dcp_world_size_and_rank(enabled: bool = True) -> tuple[int, int]:
"""Return ``(world_size, rank)`` in the DCP group, or ``(1, 0)`` when disabled
(e.g. a replicated draft cache) or the group is uninitialized (unit tests)."""
if not enabled or _DCP is None:
return 1, 0
return _DCP.world_size, _DCP.rank_in_group
_PP: GroupCoordinator | None = None
vllm/v1/attention/backend.py
+7 -8
View File
@@ -870,15 +870,14 @@ class AttentionImplBase(ABC, Generic[T]):
def __new__(cls, *args, **kwargs):
# use __new__ so that all subclasses will call this
self = super().__new__(cls)
try:
from vllm.distributed.parallel_state import get_dcp_group
from vllm.config import get_current_vllm_config_or_none
from vllm.distributed.parallel_state import get_dcp_world_size_and_rank
self.dcp_world_size = get_dcp_group().world_size
self.dcp_rank = get_dcp_group().rank_in_group
except AssertionError:
# DCP might not be initialized in testing
self.dcp_world_size = 1
self.dcp_rank = 0
# Replicated drafts run at DCP=1 inside a DCP target's process group.
config = get_current_vllm_config_or_none()
self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
config is None or config.parallel_config.decode_context_parallel_size > 1
)
try:
from vllm.distributed.parallel_state import get_pcp_group
vllm/v1/attention/backends/flash_attn.py
+7 -10
View File
@@ -60,7 +60,10 @@ from vllm.config import (
get_layers_from_vllm_config,
)
from vllm.config.cache import CacheDType
from vllm.distributed.parallel_state import get_dcp_group
from vllm.distributed.parallel_state import (
get_dcp_group,
get_dcp_world_size_and_rank,
)
from vllm.logger import init_logger
from vllm.platforms.interface import DeviceCapability
from vllm.utils.math_utils import cdiv, round_up
@@ -726,15 +729,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
num_qo_heads=self.num_heads_q,
)
try:
from vllm.distributed.parallel_state import get_dcp_group
self.dcp_world_size = get_dcp_group().world_size
self.dcp_rank = get_dcp_group().rank_in_group
except AssertionError:
# DCP might not be initialized in testing
self.dcp_world_size = 1
self.dcp_rank = 0
self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
kv_cache_spec.dcp_sharded
)
# Fused draft decode reuses the captured metadata object across draft
# steps. For DCP, build-time host-side decisions such as
vllm/v1/attention/backends/flashinfer.py
+19 -25
View File
@@ -30,7 +30,10 @@ from vllm.config import (
get_current_vllm_config_or_none,
)
from vllm.config.cache import CacheDType
from vllm.distributed.parallel_state import get_dcp_group
from vllm.distributed.parallel_state import (
get_dcp_group,
get_dcp_world_size_and_rank,
)
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
@@ -641,6 +644,7 @@ class FlashInferTrtllmAPIDecode:
@dataclass
class FlashInferMetadata:
kv_cache_layout: KVCacheLayout
num_actual_tokens: int
"""Total number of tokens in the batch (excluding padding)."""
@@ -747,17 +751,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
self._decode_cudagraph_max_bs,
self.compilation_config.max_cudagraph_capture_size,
)
try:
self.dcp_world_size = get_dcp_group().world_size
self.dcp_rank = get_dcp_group().rank_in_group
self.dcp_kv_cache_interleave_size = (
vllm_config.parallel_config.cp_kv_cache_interleave_size
)
except AssertionError:
# DCP might not be initialized in testing
self.dcp_world_size = 1
self.dcp_rank = 0
self.dcp_kv_cache_interleave_size = 1
self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
kv_cache_spec.dcp_sharded
)
self.dcp_kv_cache_interleave_size = (
vllm_config.parallel_config.cp_kv_cache_interleave_size
)
self.use_dcp = self.dcp_world_size > 1
self.dcp_a2a = (
self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
@@ -1424,6 +1423,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
# Leave prefill/decode/cascade_wrapper empty, to be populated
# case by case depending on the batch contents and backend selection.
attn_metadata = FlashInferMetadata(
kv_cache_layout=self.kv_cache_layout,
num_actual_tokens=num_actual_tokens,
slot_mapping=common_attn_metadata.slot_mapping,
q_data_type_prefill=self.q_data_type_prefill,
@@ -1904,8 +1904,6 @@ class FlashInferImpl(AttentionImpl):
num_heads, num_kv_heads, is_prefill=False
)
vllm_config = get_current_vllm_config_or_none()
# The layout is resolved after model construction, so read it lazily.
self.cache_config = vllm_config.cache_config if vllm_config else None
# Query pre-quantization needs a single dtype for the whole query tensor.
# SM90 XQA needs BF16/FP16-Q for decode and FP8 for prefill,
# so only enable this for SM100 trtllm-gen where both use FP8-Q.
@@ -1941,11 +1939,6 @@ class FlashInferImpl(AttentionImpl):
else:
self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
@property
def kv_cache_layout(self) -> KVCacheLayout:
assert self.cache_config is not None
return self.cache_config.get_resolved_kv_cache_layout()
def fused_output_quant_supported(self, quant_key: QuantKey):
if quant_key == kNvfp4Dynamic and self.is_kvcache_nvfp4:
logger.warning_once(
@@ -2111,6 +2104,7 @@ class FlashInferImpl(AttentionImpl):
# performance to make sure it does not introduce any overhead.
num_actual_tokens = attn_metadata.num_actual_tokens
kv_cache_layout = attn_metadata.kv_cache_layout
# FlashInfer treats uint8 KV cache as NVFP4. vLLM stores FP8 KV cache
# as uint8 bytes, so pass FP8 caches with their logical dtype.
@@ -2134,7 +2128,7 @@ class FlashInferImpl(AttentionImpl):
if attn_metadata.use_cascade:
# Cascade attention (rare case).
assert attn_metadata.cascade_wrapper is not None
stride_order = self.kv_cache_layout.layer_view_order
stride_order = kv_cache_layout.layer_view_order
if self.is_kvcache_nvfp4:
kv_cache_views = tuple(
cache.permute(*stride_order)
@@ -2154,7 +2148,7 @@ class FlashInferImpl(AttentionImpl):
num_decode_tokens = attn_metadata.num_decode_tokens
num_prefill_tokens = attn_metadata.num_prefill_tokens
stride_order = self.kv_cache_layout.layer_view_order
stride_order = kv_cache_layout.layer_view_order
kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
# Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1
# with TP=8). PyTorch permits non-canonical strides on size-1 dims;
@@ -2311,7 +2305,7 @@ class FlashInferImpl(AttentionImpl):
seq_lens_prefill = attn_metadata.prefill.seq_lens
# This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND"
assert get_flashinfer_layout_string(kv_cache_layout) == "HND"
assert is_strictly_contiguous(prefill_query)
assert is_strictly_contiguous(workspace_buffer)
assert is_strictly_contiguous(block_tables_prefill)
@@ -2501,7 +2495,7 @@ class FlashInferImpl(AttentionImpl):
# trtllm-gen needs HND layout on SM100. XQA is selected
# separately on SM90 and does not use this SM100 layout gate.
if decode_with_trtllm_gen:
assert get_flashinfer_layout_string(self.kv_cache_layout) == "HND"
assert get_flashinfer_layout_string(kv_cache_layout) == "HND"
else:
assert decode_with_xqa
assert is_strictly_contiguous(decode_query)
@@ -2547,7 +2541,7 @@ class FlashInferImpl(AttentionImpl):
window_left=self.window_left,
out=output_padded[:decode_query_tokens],
sinks=self.sinks,
kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
kv_layout=get_flashinfer_layout_string(kv_cache_layout),
q_len_per_req=q_len_per_req,
mask=attn_metadata.decode.mask,
q_cu_seq_lens=attn_metadata.decode.q_cu_seq_lens,
@@ -2614,7 +2608,7 @@ class FlashInferImpl(AttentionImpl):
sinks=self.sinks,
o_sf_scale=self.o_sf_scale,
out=out,
kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
kv_layout=get_flashinfer_layout_string(kv_cache_layout),
backend=attn_metadata.decode.kernel.value,
q_len_per_req=q_len_per_req,
max_q_len=max_q_len,
vllm/v1/attention/backends/utils.py
+22 -2
View File
@@ -21,7 +21,12 @@ from vllm.config.cache import _layout_from_name
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor
from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MambaSpec
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheLayout,
KVCacheSpec,
MambaSpec,
)
if TYPE_CHECKING:
from vllm.v1.core.sched.output import SchedulerOutput
@@ -286,9 +291,10 @@ def resolve_kv_cache_layout(
# specs can re-interpret HNC with different sizes as long as the total number of
# bytes is the same. If not block-compact, each spec must agree on HNC to alias
# the same page (this aliasing is done by the Hybrid Memory Allocator, HMA).
kv_cache_specs = tuple(kv_cache_specs or ())
hnc_shapes = {
(spec.num_heads, spec.num_states, spec.page_size_bytes)
for spec in kv_cache_specs or ()
for spec in kv_cache_specs
}
if len(hnc_shapes) > 1:
candidates = [m for m in candidates if m.is_block_compact]
@@ -298,6 +304,20 @@ def resolve_kv_cache_layout(
f"none is in every supported set: {supported_layouts}."
)
dcp_sharding = {
spec.dcp_sharded for spec in kv_cache_specs if isinstance(spec, AttentionSpec)
}
page_sizes = {spec.page_size_bytes for spec in kv_cache_specs}
if len(dcp_sharding) > 1 and len(page_sizes) > 1:
# Sharded target and replicated draft caches need independent groups.
# Block-outer layouts can pack those groups without equalizing pages.
candidates = [layout for layout in candidates if layout.is_block_outermost]
if not candidates:
raise ValueError(
"DCP with a replicated draft and mixed KV page sizes requires "
f"a block-outer KV cache layout; supported sets: {supported_layouts}."
)
if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None:
layout = _layout_from_name(requested)
if layout not in candidates:
vllm/v1/core/kv_cache_coordinator.py
+13 -13
View File
@@ -11,7 +11,6 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
from vllm.v1.core.kv_cache_utils import (
BlockHash,
KVCacheBlock,
dcp_world_size_for_kv_cache_spec,
)
from vllm.v1.core.single_type_kv_cache_manager import (
CrossAttentionManager,
@@ -20,6 +19,7 @@ from vllm.v1.core.single_type_kv_cache_manager import (
get_manager_for_kv_cache_spec,
)
from vllm.v1.kv_cache_interface import (
AttentionSpec,
FullAttentionSpec,
KVCacheConfig,
KVCacheSpec,
@@ -140,9 +140,7 @@ class KVCacheCoordinator(ABC):
role=kv_cache_group.role,
enable_caching=enable_caching,
kv_cache_group_id=i,
dcp_world_size=dcp_world_size_for_kv_cache_spec(
kv_cache_group.kv_cache_spec, dcp_world_size
),
dcp_world_size=dcp_world_size,
pcp_world_size=pcp_world_size,
scheduler_block_size=self.scheduler_block_size,
needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing,
@@ -664,16 +662,18 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
f"hash_block_size. block_sizes={cacheable_block_sizes}, "
f"hash_block_size={hash_block_size}"
)
assert pcp_world_size == 1, "PCP not support hybrid attn now."
if dcp_world_size > 1:
# DCP shards full-attention KV across ranks and replicates Mamba
# state; other spec types (e.g. sliding window) have no DCP-aware
# handling yet, so reject them explicitly.
for g in kv_cache_config.kv_cache_groups:
assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), (
for g in kv_cache_config.kv_cache_groups:
spec = g.kv_cache_spec
replicated = not spec.dcp_sharded
if pcp_world_size > 1:
assert isinstance(spec, FullAttentionSpec) or (
isinstance(spec, AttentionSpec) and replicated
), "PCP only supports full attention and replicated draft groups."
if dcp_world_size > 1:
assert isinstance(spec, FullAttentionSpec) or replicated, (
"DCP with hybrid KV cache layouts only supports "
"full-attention and Mamba groups, got: "
f"{type(g.kv_cache_spec).__name__}."
"full-attention, Mamba, and replicated draft groups, got: "
f"{type(spec).__name__}."
)
# Fine-grained hash hits require Mamba "align" and compatible cache
# managers in every group. TP needs hashing finer than the Mamba block;
vllm/v1/core/kv_cache_utils.py
+19 -41
View File
@@ -682,12 +682,7 @@ def hash_block_tokens(
def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
"""Return the token span of a cache block under DCP."""
layer_specs = iter_layer_specs(spec)
if len(layer_specs) > 0 and all(
isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs
):
return spec.block_size * dcp_world_size
return spec.block_size
return spec.block_size * (dcp_world_size if spec.dcp_sharded else 1)
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
return replace(spec, block_size=block_size)
def dcp_world_size_for_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> int:
"""Return the DCP size that owns this group's block geometry.
Full-attention KV (including MLA) is sharded across DCP ranks, so prefix
hashing and manager ``block_size`` use the process DCP size. Other specs
keep replicated per-rank state (Mamba, sliding window, chunked-local) and
must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated
draft group would need a different spec, not this helper.
"""
if dcp_world_size <= 1:
return 1
inner = spec
if isinstance(spec, UniformTypeKVCacheSpecs):
inner = next(iter(spec.kv_cache_specs.values()))
if isinstance(inner, FullAttentionSpec):
return dcp_world_size
return 1
def resolve_kv_cache_block_sizes(
kv_cache_config: KVCacheConfig,
vllm_config: VllmConfig,
@@ -752,6 +725,8 @@ def resolve_kv_cache_block_sizes(
groups = kv_cache_config.kv_cache_groups
if len(groups) <= 1:
if groups and not groups[0].kv_cache_spec.dcp_sharded:
dcp = 1
bs = cache_config.block_size * dcp
return bs, bs
@@ -1935,18 +1910,21 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
kv_cache_spec: The kv cache spec of each attention layer in the model
"""
if is_kv_cache_spec_uniform(
kv_cache_spec
) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):
return
logger.warning(
"Hybrid KV cache manager is disabled for this hybrid model, "
"This means we do not enable any optimizations for saving KV cache "
"memory (e.g., dropping the KV cache outside the sliding window). "
"The compute of layers like sliding window is still saved."
)
kv_cache_spec.update(_promote_local_kv_cache_specs(kv_cache_spec))
groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict)
for name, spec in kv_cache_spec.items():
replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded
groups[replicated][name] = spec
for specs in groups.values():
promoted_specs = _promote_local_kv_cache_specs(specs)
if promoted_specs == specs:
continue
logger.warning(
"Hybrid KV cache manager is disabled for this hybrid model, "
"This means we do not enable any optimizations for saving KV cache "
"memory (e.g., dropping the KV cache outside the sliding window). "
"The compute of layers like sliding window is still saved."
)
kv_cache_spec.update(promoted_specs)
def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
@@ -2276,7 +2254,7 @@ def _ensure_min_page_size(
scaled: list[KVCacheGroupSpec] = []
for g in groups:
s = g.kv_cache_spec
kw: dict[str, int] = {"block_size": s.block_size * scale}
kw: dict[str, Any] = {"block_size": s.block_size * scale}
if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None:
kw["page_size_padded"] = s.page_size_padded * scale
scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw)))
vllm/v1/core/single_type_kv_cache_manager.py
+3 -9
View File
@@ -106,11 +106,9 @@ class SingleTypeKVCacheManager(ABC):
# managers have been validated by the coordinator.
self.cache_hit_alignment_tokens = scheduler_block_size
# The block size for this manager; used for actual block allocation.
self.block_size = kv_cache_spec.block_size
self.dcp_world_size = dcp_world_size
self.pcp_world_size = pcp_world_size
if dcp_world_size > 1:
self.block_size *= dcp_world_size
self.dcp_world_size = dcp_world_size if kv_cache_spec.dcp_sharded else 1
self.pcp_world_size = pcp_world_size if kv_cache_spec.dcp_sharded else 1
self.block_size = kv_cache_spec.block_size * self.dcp_world_size
self.kv_cache_spec = kv_cache_spec
self.block_pool = block_pool
self.enable_caching = enable_caching
@@ -1455,10 +1453,6 @@ class MambaManager(SingleTypeKVCacheManager):
self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs
) -> None:
super().__init__(kv_cache_spec, block_pool, **kwargs)
# Mamba layers use TP instead of DCP, so each rank holds the full
# recurrent state. Undo the DCP/PCP block_size scaling that the base
# class applies for attention groups whose KV cache is partitioned.
self.block_size = kv_cache_spec.block_size
self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks
self.has_prefill_checkpoint_blocks = (
vllm/v1/kv_cache_interface.py
+24 -5
View File
@@ -158,6 +158,8 @@ class KVCacheSpec:
# number of tokens in a block
block_size: int
dcp_sharded: bool = field(default=False, kw_only=True)
"""Whether DCP shards this cache's token positions across ranks."""
block_stride_alignment: int | None = field(default=None, kw_only=True)
"""Required byte alignment between physical blocks, including packed layers."""
@@ -481,6 +483,7 @@ class HiSparseResidentSpec(KVCacheSpec):
@dataclass(frozen=True, kw_only=True)
class AttentionSpec(KVCacheSpec):
dcp_sharded: bool = True
num_kv_heads: int
head_size: int
dtype: torch.dtype
@@ -536,7 +539,9 @@ class AttentionSpec(KVCacheSpec):
def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
parallel_config = vllm_config.parallel_config
kv_shard_count = parallel_config.decode_context_parallel_size
kv_shard_count = (
parallel_config.decode_context_parallel_size if self.dcp_sharded else 1
)
return cdiv(max_len, self.block_size * kv_shard_count)
@@ -568,7 +573,7 @@ class FullAttentionSpec(AttentionSpec):
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
max_model_len = vllm_config.model_config.max_model_len
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
if dcp_world_size > 1:
if self.dcp_sharded and dcp_world_size > 1:
max_model_len = cdiv(max_model_len, dcp_world_size)
return cdiv(max_model_len, self.block_size) * self.page_size_bytes
@@ -612,6 +617,7 @@ class FullAttentionSpec(AttentionSpec):
head_size_v=specs[0].head_size_v,
dtype=specs[0].dtype,
kv_quant_mode=specs[0].kv_quant_mode,
dcp_sharded=specs[0].dcp_sharded,
page_size_padded=specs[0].page_size_padded,
num_head_slots=specs[0].num_head_slots,
state_content_bytes=specs[0].state_content_bytes,
@@ -698,6 +704,7 @@ class MLAAttentionSpec(FullAttentionSpec):
head_size=specs[0].head_size,
dtype=specs[0].dtype,
kv_quant_mode=specs[0].kv_quant_mode,
dcp_sharded=specs[0].dcp_sharded,
page_size_padded=specs[0].page_size_padded,
num_head_slots=specs[0].num_head_slots,
state_content_bytes=specs[0].state_content_bytes,
@@ -760,6 +767,7 @@ class RSWASpec(FullAttentionSpec):
head_size_v=base.head_size_v,
dtype=base.dtype,
kv_quant_mode=base.kv_quant_mode,
dcp_sharded=base.dcp_sharded,
page_size_padded=base.page_size_padded,
num_head_slots=base.num_head_slots,
state_content_bytes=base.state_content_bytes,
@@ -850,9 +858,10 @@ class SlidingWindowSpec(AttentionSpec):
return cdiv(num_tokens, self.block_size) + 1
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
"DCP not support sliding window."
)
assert (
not self.dcp_sharded
or vllm_config.parallel_config.decode_context_parallel_size == 1
), "DCP not support sliding window."
max_blocks = self.max_admission_blocks_per_request(
max_in_flight_tokens=vllm_config.max_in_flight_tokens,
max_model_len=vllm_config.model_config.max_model_len,
@@ -950,6 +959,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
bounded_replay_set = set(spec.bounded_replay for spec in specs)
block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
assert len({spec.dcp_sharded for spec in specs}) == 1
assert (
len(cache_dtype_str_set) == 1
and len(tokens_per_state_set) == 1
@@ -969,6 +979,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
num_kv_heads=specs[0].num_kv_heads,
head_size=specs[0].head_size,
dtype=specs[0].dtype,
dcp_sharded=specs[0].dcp_sharded,
page_size_padded=specs[0].page_size_padded,
num_head_slots=specs[0].num_head_slots,
state_content_bytes=specs[0].state_content_bytes,
@@ -1185,6 +1196,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
block_stride_alignment=specs[0].block_stride_alignment,
dtype=specs[0].dtype,
kv_quant_mode=specs[0].kv_quant_mode,
dcp_sharded=specs[0].dcp_sharded,
page_size_padded=specs[0].page_size_padded,
num_head_slots=specs[0].num_head_slots,
state_content_bytes=specs[0].state_content_bytes,
@@ -1217,6 +1229,11 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
kv_cache_specs: dict[str, KVCacheSpec]
def __post_init__(self):
super().__post_init__()
if self.kv_cache_specs:
object.__setattr__(self, "dcp_sharded", self.first_spec.dcp_sharded)
@property
def prefix_cacheable(self) -> bool:
return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
@@ -1270,6 +1287,8 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
if len(block_sizes) > 1:
# Different block sizes, not uniform.
return False
if len({spec.dcp_sharded for spec in kv_cache_specs.values()}) > 1:
return False
first_spec = next(iter(kv_cache_specs.values()))
return first_spec.is_uniform_with_collection(kv_cache_specs)
vllm/v1/simple_kv_offload/manager.py
+2 -6
View File
@@ -27,13 +27,13 @@ from vllm.v1.core.kv_cache_coordinator import (
from vllm.v1.core.kv_cache_utils import (
BlockHashWithGroupId,
ExternalBlockHash,
dcp_world_size_for_kv_cache_spec,
generate_block_hash_extra_keys,
get_block_hash,
get_group_id,
make_block_hash_with_group_id,
maybe_convert_block_hash,
resolve_block_hashes,
resolve_dcp_kv_block_size,
)
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.kv_cache_interface import (
@@ -302,11 +302,7 @@ class SimpleCPUOffloadScheduler:
target = 0
for g in kv_cache_config.prefix_cacheable_groups:
spec = g.kv_cache_spec
# Only full attention is sharded across DCP ranks; replicated specs
# (mamba, sliding window, chunked-local) keep their own block size.
block_size = spec.block_size * dcp_world_size_for_kv_cache_spec(
spec, cp_world_size
)
block_size = resolve_dcp_kv_block_size(spec, cp_world_size)
if isinstance(spec, MambaSpec):
target += 2
elif isinstance(spec, SlidingWindowSpec):
vllm/v1/spec_decode/draft_model.py
+8 -1
View File
@@ -81,13 +81,20 @@ class DraftModelProposer(SpecDecodeBaseProposer):
def _create_draft_vllm_config(self) -> VllmConfig:
base = super()._create_draft_vllm_config()
spec = self.speculative_config
target_parallel = base.parallel_config
return replace(
base,
quant_config=None,
parallel_config=replace(
spec.draft_parallel_config,
rank=self.vllm_config.parallel_config.rank,
rank=target_parallel.rank,
decode_context_parallel_size=(
target_parallel.decode_context_parallel_size
),
cp_kv_cache_interleave_size=target_parallel.cp_kv_cache_interleave_size,
dcp_comm_backend=target_parallel.dcp_comm_backend,
dcp_q_replicate=target_parallel.dcp_q_replicate,
),
model_config=spec.draft_model_config,
)
vllm/v1/worker/cp_utils.py
+2
View File
@@ -40,6 +40,8 @@ def check_attention_cp_compatibility(
layer_impl = getattr(layer, "impl", None)
if layer_impl is None:
continue
if not check_pcp and layer_impl.dcp_world_size == 1:
continue
if vllm_config.speculative_config is not None and interleave_size > 1:
assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, (
"MTP with cp_kv_cache_interleave_size > 1 is not "
vllm/v1/worker/gpu/block_table.py
+11
View File
@@ -27,6 +27,7 @@ class BlockTables:
cp_rank: int = 0,
cp_interleave: int = 1,
slot_mapping_enabled: list[bool] | None = None,
dcp_sharded: list[bool] | None = None,
):
self.block_sizes = block_sizes
self.kernel_block_sizes = kernel_block_sizes
@@ -44,6 +45,10 @@ class BlockTables:
slot_mapping_enabled = [True] * self.num_kv_cache_groups
assert len(slot_mapping_enabled) == self.num_kv_cache_groups
self._slot_mapping_enabled = slot_mapping_enabled
if dcp_sharded is None:
dcp_sharded = [True] * self.num_kv_cache_groups
assert len(dcp_sharded) == self.num_kv_cache_groups
self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device)
self.blocks_per_kv_block = [
bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
@@ -210,6 +215,7 @@ class BlockTables:
self.block_sizes_tensor,
self.kernel_block_sizes_tensor,
self.slot_mapping_enabled,
self.dcp_sharded,
slot_mappings,
slot_mappings.stride(0),
self.cp_rank,
@@ -283,6 +289,7 @@ def _compute_slot_mappings_kernel(
block_sizes, # [num_kv_cache_groups]
kernel_block_sizes, # [num_kv_cache_groups]
slot_mapping_enabled, # [num_kv_cache_groups]
dcp_sharded, # [num_kv_cache_groups]
slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
slot_mappings_stride,
cp_rank,
@@ -312,6 +319,8 @@ def _compute_slot_mappings_kernel(
kv_block_size = tl.load(block_sizes + group_id)
kernel_block_size = tl.load(kernel_block_sizes + group_id)
mapping_enabled = tl.load(slot_mapping_enabled + group_id)
if CP_SIZE != 1:
sharded = tl.load(dcp_sharded + group_id)
req_state_idx = tl.load(idx_mapping + batch_idx)
# idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns
@@ -337,6 +346,8 @@ def _compute_slot_mappings_kernel(
remainder = virtual_block_offsets % CP_INTERLEAVE
local_offsets = rounds * CP_INTERLEAVE + remainder
local_positions = virtual_block_indices * kv_block_size + local_offsets
local_positions = tl.where(sharded, local_positions, positions)
is_local = ~sharded | is_local
block_indices = tl.where(
mapping_enabled, local_positions // kernel_block_size, 0
vllm/v1/worker/gpu/model_runner.py
+14 -1
View File
@@ -21,6 +21,7 @@ import gc
import time
from contextlib import AbstractContextManager
from copy import deepcopy
from dataclasses import replace
from typing import Any, NamedTuple
import numpy as np
@@ -561,7 +562,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
return encoder_runner.get_encoder_timing_stats()
def get_kv_cache_spec(self):
return get_kv_cache_spec(self.vllm_config)
kv_cache_spec = get_kv_cache_spec(self.vllm_config)
# A draft running at DCP=1 under a DCP target keeps a replicated cache.
if (
isinstance(self.speculator, DraftModelSpeculator)
and self.speculator.dcp_size == 1
and self.dcp_size > 1
):
for name in self.speculator.draft_attn_layer_names & kv_cache_spec.keys():
kv_cache_spec[name] = replace(kv_cache_spec[name], dcp_sharded=False)
return kv_cache_spec
def initialize_kv_cache(
self,
@@ -587,6 +597,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
block_sizes = []
max_num_blocks_per_group = []
slot_mapping_enabled = []
dcp_sharded = []
for kv_cache_group in kv_cache_config.kv_cache_groups:
spec = kv_cache_group.kv_cache_spec
block_sizes.append(spec.block_size)
@@ -594,6 +605,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
)
slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
dcp_sharded.append(spec.dcp_sharded)
# Let each cache type account for CP. Attention KV is DCP-sharded,
# while Mamba/GDN recurrent state is replicated across DCP ranks.
max_num_blocks = spec.max_num_blocks_per_req(
@@ -657,6 +669,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
device=self.device,
kernel_block_sizes=self.kernel_block_sizes,
slot_mapping_enabled=slot_mapping_enabled,
dcp_sharded=dcp_sharded,
cp_size=self.dcp_size,
cp_rank=self.dcp_rank,
cp_interleave=self.cp_interleave,
vllm/v1/worker/gpu/pcp_manager.py
+4 -1
View File
@@ -145,7 +145,10 @@ class PCPManager:
if speculative_config is not None:
if speculative_config.use_dspark():
dcp_size = parallel_config.decode_context_parallel_size
if dcp_size not in (1, pcp_size):
if (
dcp_size not in (1, pcp_size)
and speculative_config.draft_model_config.use_mla
):
raise NotImplementedError(
"MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got "
f"DCP={dcp_size}, PCP={pcp_size}."
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+13 -7
View File
@@ -37,12 +37,18 @@ class DFlashSpeculator(DraftModelSpeculator):
def __init__(self, vllm_config: VllmConfig, device: torch.device):
parallel_config = vllm_config.parallel_config
if parallel_config.prefill_context_parallel_size > 1:
vllm_config = copy.copy(vllm_config)
vllm_config.parallel_config = replace(
parallel_config,
prefill_context_parallel_size=1,
)
speculative_config = vllm_config.speculative_config
assert speculative_config is not None
vllm_config = copy.copy(vllm_config)
vllm_config.parallel_config = replace(
parallel_config,
prefill_context_parallel_size=1,
decode_context_parallel_size=(
parallel_config.decode_context_parallel_size
if speculative_config.draft_model_config.use_mla
else 1
),
)
super().__init__(vllm_config, device)
self.hidden_states = torch.zeros(
@@ -427,7 +433,7 @@ class DFlashSpeculator(DraftModelSpeculator):
self.block_tables.input_block_tables[gid],
self.block_tables.kernel_block_sizes[gid],
self.block_tables.cp_rank,
self.block_tables.cp_size,
self.dcp_size,
self.block_tables.cp_interleave,
self.parallel_drafting_token_id,
self.num_query_per_req,
vllm/v1/worker/gpu/spec_decode/speculator.py
+3 -2
View File
@@ -138,6 +138,7 @@ class DraftModelSpeculator(BaseSpeculator):
# DP configuration
self.dp_size = vllm_config.parallel_config.data_parallel_size
self.dp_rank = vllm_config.parallel_config.data_parallel_rank
self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
self.eplb_state: EplbState | None = None
@@ -332,14 +333,14 @@ class DraftModelSpeculator(BaseSpeculator):
out=draft_seq_lens_cpu_upper_bound[:num_reqs],
)
draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
if dcp_local_seq_lens is None and self.dcp_size > 1:
# Draft steps advance and rewind their own global sequence lengths,
# so the target model's DCP-local lengths may already be stale.
dcp_local_seq_lens = prepare_dcp_local_seq_lens(
self.input_buffers.dcp_local_seq_lens,
self.input_buffers.seq_lens,
num_reqs,
self.block_tables.cp_size,
self.dcp_size,
self.block_tables.cp_rank,
self.block_tables.cp_interleave,
)