From 0a0d20a60a284e710a8942aaf0d238451067ad9b Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:09:45 +0000 Subject: [PATCH 1/3] Restore aligned DFlash cache geometry Assisted-by: OpenAI Codex Signed-off-by: logprobz <321553542+logprobz@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 136 ++++++++++++++++++ tests/v1/spec_decode/test_dflash_dcp.py | 55 +++++++ tests/v1/worker/test_cp_utils.py | 38 +++++ vllm/model_executor/models/qwen3_dflash.py | 36 ++++- vllm/v1/attention/backends/flash_attn.py | 16 ++- vllm/v1/core/kv_cache_coordinator.py | 31 +++- vllm/v1/core/kv_cache_utils.py | 104 +++++++++++++- vllm/v1/core/single_type_kv_cache_manager.py | 12 +- vllm/v1/kv_cache_interface.py | 37 ++++- vllm/v1/worker/cp_utils.py | 15 ++ vllm/v1/worker/gpu/block_table.py | 33 ++++- vllm/v1/worker/gpu/model_runner.py | 5 + .../gpu/spec_decode/dflash/speculator.py | 8 +- 13 files changed, 490 insertions(+), 36 deletions(-) create mode 100644 tests/v1/spec_decode/test_dflash_dcp.py diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 39f18c6d4f7c..79105604928c 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2079,6 +2079,30 @@ def test_get_kv_cache_configs_attention_free(): ] +def test_get_kv_cache_configs_preserves_model_sliding_window_retention(): + """Generic spec-decode planning must not erase model-specific retention.""" + model_config = ModelConfig(max_model_len=4096) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + sliding_window=2048, + extra_retained_tokens=2048, + ) + + configs = get_kv_cache_configs( + vllm_config, + [{"draft": spec}], + [spec.page_size_bytes * 1024], + ) + + assert configs[0].kv_cache_groups[0].kv_cache_spec.extra_retained_tokens == 2048 + + def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { @@ -2232,6 +2256,62 @@ def test_group_and_unify_kv_cache_specs_mixed_page_size_groups(): assert layer_names == {"mla.0", "mla.1", "swa.0"} +def test_group_dcp_replicated_dflash_draft(): + target = new_mla_spec() + draft = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + sliding_window=2048, + dcp_replicated=True, + ) + assert target.page_size_bytes != draft.page_size_bytes + + specs = {"model.layers.0": target, "draft.layers.0": draft} + # DeepSeek-V4's UniformType tuple planner is not needed for DFlash. + assert group_and_unify_kv_cache_specs(specs) is None + + groups = get_kv_cache_groups(_grouping_config(), specs) + draft_group = next( + group for group in groups if isinstance(group.kv_cache_spec, SlidingWindowSpec) + ) + assert all(group.kv_cache_spec.block_size == 16 for group in groups) + assert draft_group.kv_cache_spec.dcp_replicated is True + + +def test_group_dcp_replicated_dflash_with_hybrid_mla_target(): + target_full = new_mla_spec(block_size=16) + target_swa = SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + sliding_window=2048, + ) + draft = SlidingWindowSpec( + block_size=256, + num_kv_heads=4, + head_size=128, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_replicated=True, + ) + config = _grouping_config() + groups = get_kv_cache_groups( + config, + {"target.full": target_full, "target.swa": target_swa, "draft": draft}, + ) + + assert len(groups) == 3 + assert [group.layer_names for group in groups] == [ + ["target.full"], + ["target.swa"], + ["draft"], + ] + assert groups[-1].kv_cache_spec.dcp_replicated is True + + def new_indexer_mla_spec(block_size=16): # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller # page size than the main MLA attention (uint8, small head), so their pages @@ -2251,6 +2331,62 @@ def _grouping_config(): ) +def test_dflash_draft_cache_partition_is_pp1_only(): + mamba = MambaSpec( + block_size=16, + shapes=((18432,),), + dtypes=(torch.bfloat16,), + mamba_cache_mode="align", + ) + mla = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=576, + dtype=torch.float32, + ) + draft = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + sliding_window=2048, + ) + assert len({mamba.page_size_bytes, mla.page_size_bytes, draft.page_size_bytes}) == 1 + + specs = {} + for layer_index in [i for i in range(45) if i % 4 != 3]: + specs[f"model.layers.{layer_index}.linear_attn"] = mamba + for layer_index in [i for i in range(45) if i % 4 == 3]: + specs[f"model.layers.{layer_index}.self_attn"] = mla + draft_names = { + f"model.layers.{layer_index}.self_attn" for layer_index in range(45, 50) + } + for layer_name in draft_names: + specs[layer_name] = draft + + cases = ((2, 23, 11), (1, 45, 6)) + for pipeline_parallel_size, target_layers, expected_groups in cases: + config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=SimpleNamespace( + method="dflash", + attention_backend="FLASH_ATTN", + ), + model_config=SimpleNamespace( + get_num_layers=lambda parallel_config, + target_layers=target_layers: target_layers + ), + parallel_config=SimpleNamespace( + pipeline_parallel_size=pipeline_parallel_size + ), + ) + groups = get_kv_cache_groups(config, dict(specs)) + + assert len(groups) == expected_groups + if pipeline_parallel_size == 1: + assert any(set(group.layer_names) == draft_names for group in groups) + + def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): block_size = 544 full_spec = FullAttentionSpec( diff --git a/tests/v1/spec_decode/test_dflash_dcp.py b/tests/v1/spec_decode/test_dflash_dcp.py new file mode 100644 index 000000000000..01885e060704 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_dcp.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import torch + +from vllm.model_executor.models.qwen3_dflash import DFlashAttention +from vllm.v1.attention.backend import AttentionType +from vllm.v1.kv_cache_interface import SlidingWindowSpec + + +def test_dflash_sliding_window_cache_uses_aligned_block_size(): + attention = SimpleNamespace( + sliding_window=2048, + attn_type=AttentionType.DECODER, + num_kv_heads=1, + head_size=576, + head_size_v=576, + kv_cache_torch_dtype=torch.bfloat16, + kv_cache_dtype="auto", + ) + config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=2304), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + spec = DFlashAttention.get_kv_cache_spec(attention, config) + + assert isinstance(spec, SlidingWindowSpec) + assert spec.block_size == 2304 + assert spec.sliding_window == 2048 + assert spec.extra_retained_tokens == 2048 + + +def test_dflash_sliding_window_cache_is_replicated_under_dcp(): + attention = SimpleNamespace( + sliding_window=2048, + attn_type=AttentionType.DECODER, + num_kv_heads=1, + head_size=128, + head_size_v=128, + kv_cache_torch_dtype=torch.float8_e4m3fn, + kv_cache_dtype="fp8", + ) + config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16), + parallel_config=SimpleNamespace(decode_context_parallel_size=4), + ) + + spec = DFlashAttention.get_kv_cache_spec(attention, config) + + assert isinstance(spec, SlidingWindowSpec) + assert spec.sliding_window == 2048 + assert spec.extra_retained_tokens == 2048 + assert spec.dcp_replicated is True diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index b38ae4d0636c..2c960dff75d3 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch +import vllm.v1.worker.cp_utils as cp_utils from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention @@ -14,6 +17,41 @@ def test_skip_gate_only_for_zero_context(): ) +def test_replicated_draft_attention_executes_as_local_dcp(monkeypatch): + layer_impl = SimpleNamespace( + supports_mtp_with_cp_non_trivial_interleave_size=False, + need_to_return_lse_for_decode=False, + dcp_world_size=4, + dcp_rank=2, + total_cp_world_size=4, + total_cp_rank=2, + ) + layer = SimpleNamespace( + impl=layer_impl, + get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + ) + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda *_args, **_kwargs: {"draft": layer}, + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=4, + ), + speculative_config=SimpleNamespace(method="dflash"), + ) + + cp_utils.check_attention_cp_compatibility(config) + + assert layer_impl.dcp_world_size == 1 + assert layer_impl.dcp_rank == 0 + assert layer_impl.total_cp_world_size == 1 + assert layer_impl.total_cp_rank == 0 + + @pytest.mark.parametrize( "dcp_world_size,interleave_size,context_len", [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)], diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index e3db761961bc..14257e4c2e73 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses import io from collections.abc import Iterable @@ -35,6 +36,12 @@ from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + SlidingWindowSpec, + get_kv_quant_mode, +) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( get_eagle3_aux_layers_from_config, ) @@ -169,6 +176,33 @@ def _resolve_layer_attention( return sliding_window, _dflash_layer_causal(config, layer_idx) +class DFlashAttention(Attention): + """Attention whose small draft KV is replicated across DCP ranks.""" + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + dcp_replicated = vllm_config.parallel_config.decode_context_parallel_size > 1 + if self.sliding_window is not None: + assert self.attn_type == AttentionType.DECODER + return SlidingWindowSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size_v, + dtype=self.kv_cache_torch_dtype, + sliding_window=self.sliding_window, + # Prefix lookup verifies one lookahead block and then drops it. + # Keep one additional local window alive during chunked prefill + # so the proof block is not recycled before it can be hashed. + extra_retained_tokens=self.sliding_window, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + dcp_replicated=dcp_replicated, + ) + spec = super().get_kv_cache_spec(vllm_config) + if dcp_replicated and isinstance(spec, FullAttentionSpec): + spec = dataclasses.replace(spec, dcp_replicated=True) + return spec + + class DFlashQwen3Attention(nn.Module): """Attention for DFlash speculative decoding. @@ -244,7 +278,7 @@ def __init__( ) self.sliding_window = sliding_window - self.attn = Attention( + self.attn = DFlashAttention( self.num_heads, self.head_dim, self.scaling, diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 3d0244520c99..6a057ee36511 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -409,15 +409,17 @@ def __init__( self.max_num_splits = 0 # No upper bound on the number of splits. self.aot_schedule = get_flash_attn_version() == 3 - 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 + if getattr(kv_cache_spec, "dcp_replicated", False): self.dcp_world_size = 1 self.dcp_rank = 0 + else: + try: + 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 # Fused draft decode reuses the captured metadata object across draft # steps. For DCP, build-time host-side decisions such as diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 767ea0033483..87c2f257ca41 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -598,6 +598,11 @@ def __init__( # can be a multiple of hash_block_size. self.hash_block_size = hash_block_size self.dcp_world_size = dcp_world_size + self.pcp_world_size = pcp_world_size + self.has_dcp_replicated_group = any( + getattr(group.kv_cache_spec, "dcp_replicated", False) + for group in kv_cache_config.kv_cache_groups + ) group_block_sizes = [ manager.block_size for manager in self.single_type_managers ] @@ -610,14 +615,16 @@ def __init__( ) 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. + # Target attention remains DCP-sharded. Small speculative draft + # groups may instead replicate their cache and execute locally. for g in kv_cache_config.kv_cache_groups: - assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), ( + spec = g.kv_cache_spec + assert isinstance(spec, (FullAttentionSpec, MambaSpec)) or getattr( + spec, "dcp_replicated", False + ), ( "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; @@ -712,6 +719,18 @@ def verify_and_split_kv_cache_groups(self) -> None: for gid in group.group_ids: self.single_type_managers[gid].use_eagle = True + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: + if ( + self.dcp_world_size > 1 + and self.pcp_world_size == 1 + and self.has_dcp_replicated_group + ): + # Avoid enabling cascade attention for only the sharded target side + # of a target+replicated-draft hybrid. Concrete prefix replay still + # happens through find_longest_cache_hit(). + return [0] * len(self.kv_cache_config.kv_cache_groups) + return super().get_num_common_prefix_blocks(running_request_id) + def _align_cacheable(self, num_tokens: int) -> int: """Largest prefix of ``num_tokens`` a future cache hit could match. diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 337a55aa6cd0..ba74cea68006 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -15,6 +15,7 @@ from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.model_executor.models.utils import extract_layer_index from vllm.utils.hashing import xxhash, xxhash_cbor from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib @@ -670,12 +671,16 @@ def resolve_kv_cache_block_sizes( groups = kv_cache_config.kv_cache_groups if len(groups) <= 1: - bs = cache_config.block_size * dcp + dcp_replicated = len(groups) == 1 and getattr( + groups[0].kv_cache_spec, "dcp_replicated", False + ) + bs = cache_config.block_size * (1 if dcp_replicated else dcp) return bs, bs group_block_sizes = [ g.kv_cache_spec.block_size * dcp if isinstance(g.kv_cache_spec, AttentionSpec) + and not getattr(g.kv_cache_spec, "dcp_replicated", False) else g.kv_cache_spec.block_size for g in groups ] @@ -1595,6 +1600,44 @@ def group_and_unify_kv_cache_specs( return [mla_uniform_spec, *swa_uniform_specs] +def group_dcp_replicated_draft_kv_cache_specs( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Keep a replicated speculative draft separate from a sharded target. + + DFlash's small sliding-window cache has different allocation and DCP + semantics from the target cache. When both sides are independently + uniform, retain their concrete specs and native block sizes instead of + promoting or page-size-unifying the draft with the target. + """ + replicated = { + name: spec + for name, spec in kv_cache_spec.items() + if getattr(spec, "dcp_replicated", False) + } + if not replicated: + return None + sharded = { + name: spec + for name, spec in kv_cache_spec.items() + if not getattr(spec, "dcp_replicated", False) + } + if not sharded: + return None + if not is_kv_cache_spec_uniform(replicated): + return None + # The target need not itself be uniform. GLM-5.3, for example, mixes MLA + # cache layouts that the normal hybrid grouping path already understands. + # Re-enter grouping without the replicated draft so that path can preserve + # the target's native groups instead of page-unifying it with the draft. + sharded_groups = get_kv_cache_groups(vllm_config, dict(sharded)) + return [ + *sharded_groups, + *_get_kv_cache_groups_uniform_spec(replicated), + ] + + def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: """Pick a chunk size that minimizes total upward padding. @@ -1748,6 +1791,37 @@ def _largest_divisor_at_most(value: int, limit: int) -> int: return 1 +def _partition_dflash_draft_specs( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> tuple[dict[str, KVCacheSpec], dict[str, KVCacheSpec]]: + """Split appended DFlash layers from the target for PP1 grouping.""" + speculative_config = vllm_config.speculative_config + if ( + speculative_config is None + or speculative_config.method != "dflash" + or vllm_config.parallel_config.pipeline_parallel_size > 1 + ): + return kv_cache_spec, {} + + target_num_layers = vllm_config.model_config.get_num_layers( + vllm_config.parallel_config + ) + target_specs: dict[str, KVCacheSpec] = {} + draft_specs: dict[str, KVCacheSpec] = {} + for layer_name, spec in kv_cache_spec.items(): + try: + layer_index = extract_layer_index(layer_name) + except (AssertionError, IndexError, ValueError): + target_specs[layer_name] = spec + continue + if layer_index >= target_num_layers: + draft_specs[layer_name] = spec + else: + target_specs[layer_name] = spec + return target_specs, draft_specs + + def get_kv_cache_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec] ) -> list[KVCacheGroupSpec]: @@ -1769,6 +1843,19 @@ def get_kv_cache_groups( # attention free models. return [] + target_specs, draft_specs = _partition_dflash_draft_specs( + vllm_config, kv_cache_spec + ) + if target_specs and draft_specs: + target_groups = get_kv_cache_groups(vllm_config, target_specs) + draft_groups = get_kv_cache_groups(vllm_config, draft_specs) + logger.info( + "Keeping %d DFlash draft KV layers in %d independent cache groups", + len(draft_specs), + len(draft_groups), + ) + return [*target_groups, *draft_groups] + if is_kv_cache_spec_uniform(kv_cache_spec): # KV cache of all layers are the same, which is true for # most models. Allocate the same amount of memory for @@ -1779,6 +1866,10 @@ def get_kv_cache_groups( # full attention, or all layers are sliding window attention with the # same window size). Put all layers into one group. return _get_kv_cache_groups_uniform_type(uniform_spec) + elif replicated_groups := group_dcp_replicated_draft_kv_cache_specs( + vllm_config, kv_cache_spec + ): + return replicated_groups elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec): # DeepseekV4 case: All layers need the same number of token slots, # yet some layers are full attention while others are sliding window @@ -2111,9 +2202,11 @@ def get_kv_cache_configs( # This is to prevent that some layers are initialized with unregistered specs. KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) - # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # When speculating with more than 1 speculative module (e.g. multi-layered MTP), # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = ( + # A model-specific cache spec may require a larger retention window (DFlash's + # EAGLE prefix proof is one example), so never erase that value here. + mtp_extra_retained_tokens = ( vllm_config.speculative_config.num_speculative_tokens - 1 if vllm_config.speculative_config is not None and vllm_config.speculative_config.use_multi_module_mtp() @@ -2122,7 +2215,10 @@ def get_kv_cache_configs( for layer_name, layer_spec in merged_kv_cache_specs.items(): if isinstance(layer_spec, SlidingWindowSpec): merged_kv_cache_specs[layer_name] = replace( - layer_spec, extra_retained_tokens=extra_retained_tokens + layer_spec, + extra_retained_tokens=max( + layer_spec.extra_retained_tokens, mtp_extra_retained_tokens + ), ) # Get global KV cache groups. This also handles spec unification for diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index ceb56a59875f..6ea391dc617f 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -75,7 +75,7 @@ def __init__( 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: + if dcp_world_size > 1 and not getattr(kv_cache_spec, "dcp_replicated", False): self.block_size *= dcp_world_size self.kv_cache_spec = kv_cache_spec self.block_pool = block_pool @@ -700,7 +700,7 @@ def find_longest_cache_hit( "and chunked local attention groups" ) block_size = kv_cache_spec.block_size - if dcp_world_size > 1: + if dcp_world_size > 1 and not getattr(kv_cache_spec, "dcp_replicated", False): # DCP shards each block's KV across ranks; hashes must be viewed at # the sharded block size. block_size *= dcp_world_size @@ -915,8 +915,12 @@ def find_longest_cache_hit( assert isinstance(kv_cache_spec, SlidingWindowSpec), ( "SlidingWindowManager can only be used for sliding window groups" ) - assert dcp_world_size == 1, "DCP not support sliding window attn now." - assert pcp_world_size == 1, "PCP not support sliding window attn now." + assert dcp_world_size == 1 or kv_cache_spec.dcp_replicated, ( + "DCP only supports sliding-window KV when it is replicated." + ) + assert pcp_world_size == 1 or kv_cache_spec.dcp_replicated, ( + "PCP only supports sliding-window KV when it is replicated." + ) # Fine-grained partial hits are not supported for sliding window now assert alignment_tokens % kv_cache_spec.block_size == 0, ( "SlidingWindowManager does not support fine-grained (partial) cache hits" diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 3cecb4fbee8a..8f9483960136 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -228,8 +228,11 @@ def is_uniform_with_collection( f"Unsupported KV cache spec type: {type(self)}. " "Please register it using @register_kv_cache_spec decorator." ) + dcp_replicated = getattr(self, "dcp_replicated", False) return all( - isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() + isinstance(spec, uniform_type_base_spec) + and getattr(spec, "dcp_replicated", False) == dcp_replicated + for spec in kv_cache_specs.values() ) @@ -428,7 +431,11 @@ def real_page_size_bytes(self) -> int: 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 = ( + 1 + if getattr(self, "dcp_replicated", False) + else parallel_config.decode_context_parallel_size + ) return cdiv(max_len, self.block_size * kv_shard_count) @@ -458,10 +465,12 @@ class FullAttentionSpec(AttentionSpec): cache layout itself. """ + dcp_replicated: bool = False + 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 dcp_world_size > 1 and not self.dcp_replicated: max_model_len = cdiv(max_model_len, dcp_world_size) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @@ -498,6 +507,11 @@ def merge(cls, specs: list[Self]) -> Self: assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), ( "MLAAttentionSpec should be merged in MLAAttentionSpec.merge" ) + dcp_replicated = {spec.dcp_replicated for spec in specs} + assert len(dcp_replicated) == 1, ( + "All attention layers in one KV cache group must use the same " + "DCP replication mode." + ) merged_spec = cls( block_size=specs[0].block_size, num_kv_heads=specs[0].num_kv_heads, @@ -514,6 +528,7 @@ def merge(cls, specs: list[Self]) -> Self: # If any layer in the group is non-causal, treat the group as # non-causal so the engine core disables incompatible scheduling. non_causal=any(spec.non_causal for spec in specs), + dcp_replicated=dcp_replicated.pop(), ) for spec in specs: for f in fields(AttentionSpec): @@ -703,6 +718,7 @@ def is_uniform_with_collection( @dataclass(frozen=True, kw_only=True) class SlidingWindowSpec(AttentionSpec): sliding_window: int + dcp_replicated: bool = False # The trailing edge of the window is extended by ``extra_retained_tokens`` # so that those extra trailing tokens' blocks are retained (but not # attended). This is needed for multi-module spec decoding which can @@ -739,9 +755,10 @@ def max_admission_blocks_per_request( 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 ( + vllm_config.parallel_config.decode_context_parallel_size == 1 + or self.dcp_replicated + ), "DCP only supports sliding-window KV when it is replicated." 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, @@ -754,6 +771,7 @@ def is_uniform_with_collection( return all( isinstance(spec, SlidingWindowSpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -980,6 +998,13 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): def page_size_bytes(self) -> int: return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) + @property + def dcp_replicated(self) -> bool: + return all( + getattr(spec, "dcp_replicated", False) + for spec in self.kv_cache_specs.values() + ) + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_num_pages = max( cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 92d8383c1f12..8d7d272f7680 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -37,6 +37,21 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: layer_impl = getattr(layer, "impl", None) if layer_impl is None: continue + get_spec = getattr(layer, "get_kv_cache_spec", None) + if get_spec is not None: + try: + spec = get_spec(vllm_config) + except Exception: + spec = None + if getattr(spec, "dcp_replicated", False): + # Replicated draft KV contains the complete sequence on + # every rank, so its attention executes as a local DCP1 op. + layer_impl.dcp_world_size = 1 + layer_impl.dcp_rank = 0 + layer_impl.total_cp_world_size = 1 + layer_impl.total_cp_rank = 0 + layer_impl.need_to_return_lse_for_decode = False + 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 " diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index dae1630a2c19..2a70cedffc0e 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -26,6 +26,7 @@ def __init__( cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, + group_cp_sizes: list[int] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -36,6 +37,12 @@ def __init__( self.cp_size = cp_size self.cp_rank = cp_rank self.cp_interleave = cp_interleave + if group_cp_sizes is None: + group_cp_sizes = [cp_size] * len(block_sizes) + assert len(group_cp_sizes) == len(block_sizes) + self.group_cp_sizes = torch.tensor( + group_cp_sizes, dtype=torch.int32, device=device + ) self.num_kv_cache_groups = len(self.block_sizes) assert len(max_num_blocks_per_group) == self.num_kv_cache_groups @@ -205,8 +212,11 @@ def compute_slot_mappings( positions, self.block_table_ptrs, self.block_table_strides, + self.num_blocks.gpu, + self.num_blocks.gpu.stride(0), self.block_sizes_tensor, self.kernel_block_sizes_tensor, + self.group_cp_sizes, slot_mappings, slot_mappings.stride(0), self.cp_rank, @@ -268,6 +278,10 @@ def _gather_block_tables_kernel( block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks) tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks) + for i in tl.range(num_blocks, max_num_blocks, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + tl.store(dst_row_ptr + offset, 0, mask=offset < max_num_blocks) + @triton.jit def _compute_slot_mappings_kernel( @@ -277,8 +291,11 @@ def _compute_slot_mappings_kernel( pos, # [num_tokens] block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] + num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs] + num_blocks_stride, block_sizes, # [num_kv_cache_groups] kernel_block_sizes, # [num_kv_cache_groups] + group_cp_sizes, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, cp_rank, @@ -305,20 +322,24 @@ def _compute_slot_mappings_kernel( block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) block_table_stride = tl.load(block_table_strides + group_id) + group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride kv_block_size = tl.load(block_sizes + group_id) kernel_block_size = tl.load(kernel_block_sizes + group_id) + group_cp_size = tl.load(group_cp_sizes + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) + num_blocks = tl.load(group_num_blocks_ptr + req_state_idx) start_idx = tl.load(query_start_loc + batch_idx) end_idx = tl.load(query_start_loc + batch_idx + 1) for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) - positions = tl.load(pos + offset, mask=offset < end_idx, other=0) + token_mask = offset < end_idx + positions = tl.load(pos + offset, mask=token_mask, other=0) - if CP_SIZE == 1: + if CP_SIZE == 1 or group_cp_size == 1: # Common case: Context parallelism is not used. local_positions = positions - is_local = True + is_local = token_mask else: # Context parallelism is used. virtual_block_size = kv_block_size * CP_SIZE @@ -332,13 +353,15 @@ def _compute_slot_mappings_kernel( block_indices = local_positions // kernel_block_size block_offsets = local_positions % kernel_block_size + valid_block = token_mask & (block_indices < num_blocks) block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, - mask=is_local, + mask=is_local & valid_block, other=0, ) slot_ids = block_numbers * kernel_block_size + block_offsets - if CP_SIZE != 1: + if CP_SIZE != 1 and group_cp_size != 1: slot_ids = tl.where(is_local, slot_ids, PAD_ID) + slot_ids = tl.where(valid_block, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 049f512e19f0..aa8ff7677fd5 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -551,9 +551,13 @@ def initialize_kv_cache( block_sizes = [] max_num_blocks_per_group = [] + group_cp_sizes = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) + group_cp_sizes.append( + 1 if getattr(spec, "dcp_replicated", False) else self.dcp_size + ) # 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( @@ -611,6 +615,7 @@ def initialize_kv_cache( cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, + group_cp_sizes=group_cp_sizes, ) self.pcp_manager = pcp.maybe_build_pcp_manager( self.vllm_config, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 651d5905c3b3..dd8c32b4ceed 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -421,9 +421,11 @@ def propose( seeds, 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.block_tables.cp_interleave, + # Every DFlash draft cache group is replicated under DCP, so + # draft context/query slots are ordinary local DCP1 slots. + 0, + 1, + 1, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, From dfd69b27123be7a850aa4f1e86481cf85391af13 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Sun, 30 Aug 2026 08:20:37 +0000 Subject: [PATCH 2/3] Harden replicated DFlash cache handling Assisted-by: OpenAI Codex Signed-off-by: logprobz <321553542+logprobz@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 98 ++++++++++++++++ tests/v1/spec_decode/test_dflash_dcp.py | 103 ++++++++++++++++- tests/v1/worker/test_cp_utils.py | 72 ++++++++++++ tests/v1/worker/test_gpu_block_table.py | 105 ++++++++++++++++++ vllm/model_executor/models/qwen3_dflash.py | 3 + vllm/v1/attention/backend.py | 4 + vllm/v1/attention/backends/flash_attn.py | 1 + vllm/v1/core/kv_cache_utils.py | 1 + vllm/v1/kv_cache_interface.py | 24 +++- vllm/v1/worker/cp_utils.py | 13 ++- vllm/v1/worker/gpu/block_table.py | 14 ++- .../gpu/spec_decode/dflash/speculator.py | 33 ++++-- 12 files changed, 447 insertions(+), 24 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 79105604928c..6930057fd92f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -55,6 +55,7 @@ KVQuantMode, MambaSpec, MLAAttentionSpec, + RSWASpec, SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, @@ -2331,6 +2332,103 @@ def _grouping_config(): ) +@pytest.mark.parametrize( + "spec", + [ + MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + dcp_replicated=True, + ), + RSWASpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + dcp_replicated=True, + rswa_window=64, + ), + SinkFullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + dcp_replicated=True, + sink_len=4, + ), + SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + dcp_replicated=True, + sliding_window=64, + ), + ], + ids=("mla", "rswa", "sink", "sliding-window-mla"), +) +def test_attention_spec_merge_preserves_dcp_replicated(spec): + merged = type(spec).merge([spec, spec]) + + assert merged.dcp_replicated is True + + +def test_sliding_window_mla_uniformity_includes_dcp_replication(): + replicated = SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + dcp_replicated=True, + sliding_window=64, + ) + sharded = SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + dcp_replicated=False, + sliding_window=64, + ) + + assert not replicated.is_uniform_with_collection( + {"replicated": replicated, "sharded": sharded} + ) + + +def test_disable_hybrid_manager_skips_dflash_partition(): + target = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + ) + draft = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + sliding_window=64, + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=True), + speculative_config=SimpleNamespace(method="dflash"), + model_config=SimpleNamespace(get_num_layers=lambda _parallel_config: 1), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + ) + specs = { + "model.layers.0.self_attn": target, + "model.layers.1.self_attn": draft, + } + + groups = get_kv_cache_groups(config, specs) + + assert len(groups) == 1 + assert set(groups[0].layer_names) == set(specs) + + def test_dflash_draft_cache_partition_is_pp1_only(): mamba = MambaSpec( block_size=16, diff --git a/tests/v1/spec_decode/test_dflash_dcp.py b/tests/v1/spec_decode/test_dflash_dcp.py index 01885e060704..053a36ef415d 100644 --- a/tests/v1/spec_decode/test_dflash_dcp.py +++ b/tests/v1/spec_decode/test_dflash_dcp.py @@ -4,6 +4,8 @@ import torch +import vllm.v1.worker.gpu.spec_decode.dflash.speculator as dflash_speculator +from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.models.qwen3_dflash import DFlashAttention from vllm.v1.attention.backend import AttentionType from vllm.v1.kv_cache_interface import SlidingWindowSpec @@ -19,8 +21,12 @@ def test_dflash_sliding_window_cache_uses_aligned_block_size(): kv_cache_torch_dtype=torch.bfloat16, kv_cache_dtype="auto", ) + padded_page_size = 6 * 1024 * 1024 config = SimpleNamespace( - cache_config=SimpleNamespace(block_size=2304), + cache_config=SimpleNamespace( + block_size=2304, + skip_page_size_padded=padded_page_size, + ), parallel_config=SimpleNamespace(decode_context_parallel_size=1), ) @@ -30,6 +36,7 @@ def test_dflash_sliding_window_cache_uses_aligned_block_size(): assert spec.block_size == 2304 assert spec.sliding_window == 2048 assert spec.extra_retained_tokens == 2048 + assert spec.page_size_padded == padded_page_size def test_dflash_sliding_window_cache_is_replicated_under_dcp(): @@ -53,3 +60,97 @@ def test_dflash_sliding_window_cache_is_replicated_under_dcp(): assert spec.sliding_window == 2048 assert spec.extra_retained_tokens == 2048 assert spec.dcp_replicated is True + + +def test_dflash_uses_draft_group_dcp_slot_parameters(monkeypatch): + """A sharded DSpark-style draft group must retain the real DCP mapping.""" + cp_args = [] + + def capture_prepare_args(*args): + cp_args.append(args[18:21]) + + monkeypatch.setattr( + dflash_speculator, "prepare_dflash_inputs", capture_prepare_args + ) + monkeypatch.setattr( + dflash_speculator, + "dispatch_cg_and_sync_dp", + lambda *_args, **_kwargs: ( + SimpleNamespace( + num_reqs=1, + num_tokens=1, + cg_mode=CUDAGraphMode.NONE, + ), + None, + ), + ) + monkeypatch.setattr( + dflash_speculator, + "build_slot_mappings_by_layer", + lambda *_args, **_kwargs: {}, + ) + + block_tables = SimpleNamespace( + slot_mappings=torch.zeros((1, 1), dtype=torch.int64), + input_block_tables=[torch.zeros((1, 1), dtype=torch.int32)], + kernel_block_sizes=[16], + get_group_cp_parameters=lambda _gid: (2, 4, 8), + ) + model = SimpleNamespace(precompute_and_store_context_kv=lambda *_args: None) + speculator = SimpleNamespace( + num_query_per_req=1, + num_speculative_steps=1, + max_model_len=128, + max_num_reqs=1, + max_num_tokens=1, + hidden_states=torch.zeros((1, 1)), + context_positions=torch.zeros(1, dtype=torch.int64), + sample_indices=torch.zeros(1, dtype=torch.int64), + sample_pos=torch.zeros(1, dtype=torch.int64), + sample_idx_mapping=torch.zeros(1, dtype=torch.int32), + temperature=torch.zeros(1), + seeds=torch.zeros(1, dtype=torch.int64), + input_buffers=SimpleNamespace(), + block_tables=block_tables, + draft_kv_cache_group_id=0, + draft_kv_cache_group_ids=[0], + _context_slot_mappings=torch.zeros((1, 1), dtype=torch.int64), + _layer_group_idx=None, + parallel_drafting_token_id=0, + sample_from_anchor=False, + model=model, + query_cudagraph_manager=None, + dp_size=1, + dp_rank=0, + _group_causal=False, + kv_cache_config=SimpleNamespace(), + draft_tokens=torch.zeros((1, 1), dtype=torch.int64), + _build_draft_attn_metadata=lambda **_kwargs: {}, + _prepare_eplb_forward=lambda *_args: None, + _generate_draft=lambda *_args, **_kwargs: None, + ) + input_batch = SimpleNamespace( + num_reqs=1, + num_tokens=1, + seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32), + ) + one_i32 = torch.zeros(1, dtype=torch.int32) + one_i64 = torch.zeros(1, dtype=torch.int64) + one_f32 = torch.zeros(1) + + dflash_speculator.DFlashSpeculator.propose( + speculator, + input_batch=input_batch, + attn_metadata={}, + slot_mappings={}, + last_hidden_states=torch.zeros((1, 1)), + aux_hidden_states=None, + num_sampled=one_i32, + num_rejected=one_i32, + last_sampled=one_i64, + next_prefill_tokens=one_i64, + temperature=one_f32, + seeds=one_i64, + ) + + assert cp_args == [(2, 4, 8)] diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index 2c960dff75d3..4e67c79494e8 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -26,9 +26,14 @@ def test_replicated_draft_attention_executes_as_local_dcp(monkeypatch): total_cp_world_size=4, total_cp_rank=2, ) + backend = SimpleNamespace( + supports_dcp_replicated=True, + get_name=lambda: "TEST_SUPPORTED", + ) layer = SimpleNamespace( impl=layer_impl, get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + get_attn_backend=lambda: backend, ) monkeypatch.setattr( cp_utils, @@ -52,6 +57,73 @@ def test_replicated_draft_attention_executes_as_local_dcp(monkeypatch): assert layer_impl.total_cp_rank == 0 +def test_replicated_draft_rejects_backend_without_local_dcp(monkeypatch): + layer_impl = SimpleNamespace( + supports_mtp_with_cp_non_trivial_interleave_size=False, + need_to_return_lse_for_decode=False, + ) + backend = SimpleNamespace( + supports_dcp_replicated=False, + get_name=lambda: "TEST_UNSUPPORTED", + ) + layer = SimpleNamespace( + impl=layer_impl, + get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + get_attn_backend=lambda: backend, + ) + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda *_args, **_kwargs: {"draft": layer}, + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ), + speculative_config=SimpleNamespace(method="dflash"), + ) + + with pytest.raises(AssertionError, match="replicated DCP"): + cp_utils.check_attention_cp_compatibility(config) + + +def test_attention_cache_spec_errors_are_not_swallowed(monkeypatch): + layer_impl = SimpleNamespace( + supports_mtp_with_cp_non_trivial_interleave_size=False, + need_to_return_lse_for_decode=True, + ) + + def raise_spec_error(_config): + raise RuntimeError("invalid cache spec") + + layer = SimpleNamespace( + impl=layer_impl, + get_kv_cache_spec=raise_spec_error, + get_attn_backend=lambda: SimpleNamespace( + supports_dcp_replicated=False, + get_name=lambda: "TEST", + ), + ) + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda *_args, **_kwargs: {"layer": layer}, + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ), + speculative_config=None, + ) + + with pytest.raises(RuntimeError, match="invalid cache spec"): + cp_utils.check_attention_cp_compatibility(config) + + @pytest.mark.parametrize( "dcp_world_size,interleave_size,context_len", [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)], diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index ee44ff24d581..3efac2a00bdf 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -176,6 +176,111 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): assert torch.equal(actual, expected) +def test_mixed_group_cp_slot_mapping(): + """Replicated draft and sharded target groups use independent CP geometry.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[128, 128], + max_num_reqs=1, + max_num_batched_tokens=1024, + max_num_blocks_per_group=[2, 8], + device=device, + kernel_block_sizes=[128, 128], + cp_size=4, + cp_rank=1, + cp_interleave=128, + group_cp_sizes=[4, 1], + ) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([5, 9], list(range(20, 28))), + overwrite=True, + ) + block_tables.apply_staged_writes() + + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 1024], dtype=torch.int32, device=device) + positions = torch.arange(1024, dtype=torch.int64, device=device) + actual = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=1024, + ) + + expected_sharded = torch.full((1024,), -1, dtype=torch.int64, device=device) + expected_sharded[128:256] = torch.arange( + 5 * 128, 6 * 128, dtype=torch.int64, device=device + ) + expected_sharded[640:768] = torch.arange( + 9 * 128, 10 * 128, dtype=torch.int64, device=device + ) + expected_replicated = torch.cat( + [ + torch.arange(block_id * 128, (block_id + 1) * 128, device=device) + for block_id in range(20, 28) + ] + ) + + assert torch.equal(actual[0], expected_sharded) + assert torch.equal(actual[1], expected_replicated) + + +def test_group_cp_sizes_rebuilt_with_layout_tensors(): + """CuMem wake-up must restore group CP constants, not undefined storage.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[128, 128], + max_num_reqs=1, + max_num_batched_tokens=16, + max_num_blocks_per_group=[1, 1], + device=device, + kernel_block_sizes=[128, 128], + cp_size=4, + cp_rank=1, + cp_interleave=128, + group_cp_sizes=[4, 1], + ) + block_tables.group_cp_sizes.fill_(99) + + block_tables.init_block_table_layout_tensors() + + assert block_tables.group_cp_sizes.tolist() == [4, 1] + + +def test_gather_block_tables_clears_shortened_row_tail(): + """A shortened request row must not retain block ids from its prior owner.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16], + max_num_reqs=1, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[16], + ) + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([7, 8, 9],), + overwrite=True, + ) + block_tables.apply_staged_writes() + block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([11],), + overwrite=True, + ) + block_tables.apply_staged_writes() + + gathered = block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1)[0] + torch.accelerator.synchronize() + + assert gathered[0, 0].item() == 11 + assert (gathered[0, 1:] == 0).all() + + def test_v1_block_table_move_row_clears_vacated_row(): """condense() moves the last row into a freed slot; the vacated row must not keep stale block ids. Padded dummy-run batches dereference stale rows diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 14257e4c2e73..1147ba9aaa72 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -190,6 +190,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: head_size_v=self.head_size_v, dtype=self.kv_cache_torch_dtype, sliding_window=self.sliding_window, + page_size_padded=getattr( + vllm_config.cache_config, "skip_page_size_padded", None + ), # Prefix lookup verifies one lookahead block and then drops it. # Keep one additional local window alive during chunked prefill # so the proof block is not recycled before it can be hashed. diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 4fdfafaaea9a..17f8c31ee3df 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -69,6 +69,10 @@ class AttentionBackend(ABC): # Does attention's forward() include kv cache update? forward_includes_kv_cache_update: bool = True + # Whether metadata builders and kernels can execute a DCP-replicated cache + # group as a local DCP1 operation inside a larger DCP world. + supports_dcp_replicated: ClassVar[bool] = False + @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(1)] diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 6a057ee36511..42bbbd922a64 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -91,6 +91,7 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(16)] forward_includes_kv_cache_update: bool = False + supports_dcp_replicated: ClassVar[bool] = True @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index ba74cea68006..4ed3e26abd93 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1801,6 +1801,7 @@ def _partition_dflash_draft_specs( speculative_config is None or speculative_config.method != "dflash" or vllm_config.parallel_config.pipeline_parallel_size > 1 + or vllm_config.scheduler_config.disable_hybrid_kv_cache_manager ): return kv_cache_spec, {} diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8f9483960136..49ec9dce4ca0 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -393,6 +393,8 @@ class AttentionSpec(KVCacheSpec): """Tokens covered by one stored state. Ints > 1 compress multiple tokens into one state (DSv4 sparse MLA); fractions < 1 store multiple states per token (Whisper block pooling: ``Fraction(1, block_pool_size)``).""" + dcp_replicated: bool = False + """Whether every DCP rank stores the complete KV cache for this layer.""" def __post_init__(self): if self.head_size_v is None: @@ -465,8 +467,6 @@ class FullAttentionSpec(AttentionSpec): cache layout itself. """ - dcp_replicated: bool = False - 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 @@ -591,14 +591,17 @@ def merge(cls, specs: list[Self]) -> Self: page_tail_bytes_per_token_set = set( spec.page_tail_bytes_per_token for spec in specs ) + dcp_replicated_set = {spec.dcp_replicated for spec in specs} assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 and len(page_tail_bytes_per_token_set) == 1 + and len(dcp_replicated_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, tokens per state, model version, and page tail." + "quantization method, tokens per state, model version, page tail, " + "and DCP replication mode." ) merged_spec = cls( block_size=specs[0].block_size, @@ -616,6 +619,7 @@ def merge(cls, specs: list[Self]) -> Self: non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs ), + dcp_replicated=dcp_replicated_set.pop(), ) for spec in specs: for f in fields(AttentionSpec): @@ -671,6 +675,7 @@ def merge(cls, specs: list[RSWASpec]) -> RSWASpec: sliding_window=base.sliding_window, attention_chunk_size=base.attention_chunk_size, non_causal=base.non_causal, + dcp_replicated=base.dcp_replicated, rswa_window=rswa_windows.pop(), ) @@ -718,7 +723,6 @@ def is_uniform_with_collection( @dataclass(frozen=True, kw_only=True) class SlidingWindowSpec(AttentionSpec): sliding_window: int - dcp_replicated: bool = False # The trailing edge of the window is extended by ``extra_retained_tokens`` # so that those extra trailing tokens' blocks are retained (but not # attended). This is needed for multi-module spec decoding which can @@ -806,16 +810,18 @@ def merge(cls, specs: list[Self]) -> Self: model_version_set = set(spec.model_version for spec in specs) sliding_window_set = set(spec.sliding_window for spec in specs) extra_retained_set = set(spec.extra_retained_tokens for spec in specs) + dcp_replicated_set = {spec.dcp_replicated for spec in specs} assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 and len(sliding_window_set) == 1 and len(extra_retained_set) == 1 + and len(dcp_replicated_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, sliding " - "window size, and retained token count." + "window size, retained token count, and DCP replication mode." ) return cls( block_size=specs[0].block_size, @@ -830,6 +836,7 @@ def merge(cls, specs: list[Self]) -> Self: cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), + dcp_replicated=dcp_replicated_set.pop(), ) def is_uniform_with_collection( @@ -838,6 +845,7 @@ def is_uniform_with_collection( return all( isinstance(spec, SlidingWindowMLASpec) and spec.sliding_window == self.sliding_window + and spec.dcp_replicated == self.dcp_replicated for spec in kv_cache_specs.values() ) @@ -953,6 +961,11 @@ def merge(cls, specs: list[Self]) -> Self: assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), ( "MLAAttentionSpec should be merged in MLAAttentionSpec.merge" ) + dcp_replicated = {spec.dcp_replicated for spec in specs} + assert len(dcp_replicated) == 1, ( + "All attention layers in one KV cache group must use the same " + "DCP replication mode." + ) merged_spec = cls( block_size=specs[0].block_size, num_kv_heads=specs[0].num_kv_heads, @@ -967,6 +980,7 @@ def merge(cls, specs: list[Self]) -> Self: sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), non_causal=any(spec.non_causal for spec in specs), + dcp_replicated=dcp_replicated.pop(), ) for spec in specs: for f in fields(AttentionSpec): diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 8d7d272f7680..a078f3324fb5 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -28,8 +28,8 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: layers = get_layers_from_vllm_config(vllm_config, layer_type) for layer in layers.values(): get_attn_backend = getattr(layer, "get_attn_backend", None) - if pcp_size > 1 and get_attn_backend is not None: - backend = get_attn_backend() + backend = get_attn_backend() if get_attn_backend is not None else None + if pcp_size > 1 and backend is not None: assert backend.supports_pcp(), ( "PCP requires attention backend support, " f"but {backend.get_name()} does not support PCP." @@ -39,11 +39,12 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: continue get_spec = getattr(layer, "get_kv_cache_spec", None) if get_spec is not None: - try: - spec = get_spec(vllm_config) - except Exception: - spec = None + spec = get_spec(vllm_config) if getattr(spec, "dcp_replicated", False): + assert backend is None or backend.supports_dcp_replicated, ( + "Attention with replicated DCP requires backend support, " + f"but {backend.get_name()} does not provide it." + ) # Replicated draft KV contains the complete sequence on # every rank, so its attention executes as a local DCP1 op. layer_impl.dcp_world_size = 1 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 2a70cedffc0e..60dd809170af 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -40,9 +40,8 @@ def __init__( if group_cp_sizes is None: group_cp_sizes = [cp_size] * len(block_sizes) assert len(group_cp_sizes) == len(block_sizes) - self.group_cp_sizes = torch.tensor( - group_cp_sizes, dtype=torch.int32, device=device - ) + assert all(group_cp_size in (1, cp_size) for group_cp_size in group_cp_sizes) + self.group_cp_sizes_list = list(group_cp_sizes) self.num_kv_cache_groups = len(self.block_sizes) assert len(max_num_blocks_per_group) == self.num_kv_cache_groups @@ -112,8 +111,17 @@ def init_block_table_layout_tensors(self) -> None: self.kernel_block_sizes_tensor = torch.tensor( self.kernel_block_sizes, dtype=torch.int32, device=self.device ) + self.group_cp_sizes = torch.tensor( + self.group_cp_sizes_list, dtype=torch.int32, device=self.device + ) self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) + def get_group_cp_parameters(self, group_id: int) -> tuple[int, int, int]: + group_cp_size = self.group_cp_sizes_list[group_id] + if group_cp_size == 1: + return 0, 1, 1 + return self.cp_rank, group_cp_size, self.cp_interleave + def append_block_ids( self, req_index: int, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index dd8c32b4ceed..27e0754bfc7f 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -183,6 +183,17 @@ def set_attn( ] assert self.draft_kv_cache_group_ids, "No draft attention groups found." self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] + draft_cp_parameters = { + block_tables.get_group_cp_parameters(gid) + for gid in self.draft_kv_cache_group_ids + } + if len(draft_cp_parameters) != 1: + raise NotImplementedError( + "DFlash draft attention groups must use one DCP layout." + ) + self.draft_cp_rank, self.draft_cp_size, self.draft_cp_interleave = ( + draft_cp_parameters.pop() + ) # Per-group context slot buffers for the precompute (one row per group). self._context_slot_mappings = torch.zeros( @@ -222,6 +233,9 @@ def reset_attn(self) -> None: for name in ( "draft_kv_cache_group_ids", "draft_kv_cache_group_id", + "draft_cp_rank", + "draft_cp_size", + "draft_cp_interleave", "_context_slot_mappings", "_layer_group_idx", "_group_causal", @@ -304,14 +318,14 @@ def _build_draft_attn_metadata( if not self.draft_attn_layer_names: return None assert num_query_per_req is None # Omitted for DFlash, read from self instead - if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: + if dcp_local_seq_lens is None and self.draft_cp_size > 1: 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.block_tables.cp_rank, - self.block_tables.cp_interleave, + self.draft_cp_size, + self.draft_cp_rank, + self.draft_cp_interleave, ) dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens return super()._build_draft_attn_metadata( @@ -402,6 +416,9 @@ def propose( assert self.draft_kv_cache_group_id >= 0 # Support multiple draft KV cache groups by preparing inputs once for each for i, gid in enumerate(self.draft_kv_cache_group_ids): + cp_rank, cp_size, cp_interleave = self.block_tables.get_group_cp_parameters( + gid + ) prepare_dflash_inputs( self.input_buffers, self.block_tables.slot_mappings[gid], @@ -421,11 +438,9 @@ def propose( seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], - # Every DFlash draft cache group is replicated under DCP, so - # draft context/query slots are ordinary local DCP1 slots. - 0, - 1, - 1, + cp_rank, + cp_size, + cp_interleave, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, From 938d184a66369359a81e3264ab213e7d42540d44 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Sun, 30 Aug 2026 09:52:05 +0000 Subject: [PATCH 3/3] Reject incompatible replicated DCP cache groups Co-authored-by: OpenAI Codex Signed-off-by: logprobz <321553542+logprobz@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 32 +++++++++++++++ tests/v1/spec_decode/test_dflash_dcp.py | 40 +++++++++++++++++++ tests/v1/worker/test_cp_utils.py | 26 ++++++++++++ vllm/v1/core/kv_cache_utils.py | 12 ++++++ vllm/v1/worker/cp_utils.py | 6 ++- .../gpu/spec_decode/dflash/speculator.py | 23 ++++++----- 6 files changed, 128 insertions(+), 11 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 6930057fd92f..3e96546faa4f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2429,6 +2429,38 @@ def test_disable_hybrid_manager_skips_dflash_partition(): assert set(groups[0].layer_names) == set(specs) +def test_disable_hybrid_manager_rejects_mixed_dcp_replication(): + target = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + ) + draft = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + sliding_window=64, + dcp_replicated=True, + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=True), + speculative_config=SimpleNamespace(method="dflash"), + model_config=SimpleNamespace(get_num_layers=lambda _parallel_config: 1), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + ) + + with pytest.raises(ValueError, match="DCP-replicated and DCP-sharded"): + get_kv_cache_groups( + config, + { + "model.layers.0.self_attn": target, + "model.layers.1.self_attn": draft, + }, + ) + + def test_dflash_draft_cache_partition_is_pp1_only(): mamba = MambaSpec( block_size=16, diff --git a/tests/v1/spec_decode/test_dflash_dcp.py b/tests/v1/spec_decode/test_dflash_dcp.py index 053a36ef415d..cf3bcb2ac6b3 100644 --- a/tests/v1/spec_decode/test_dflash_dcp.py +++ b/tests/v1/spec_decode/test_dflash_dcp.py @@ -154,3 +154,43 @@ def capture_prepare_args(*args): ) assert cp_args == [(2, 4, 8)] + + +def test_replicated_draft_metadata_uses_full_sequence_lengths(monkeypatch): + seq_lens = torch.tensor([128, 64], dtype=torch.int32) + captured = [] + + def capture_metadata(_self, *_args, **kwargs): + captured.append(kwargs["dcp_local_seq_lens"]) + return {} + + monkeypatch.setattr( + dflash_speculator.DraftModelSpeculator, + "_build_draft_attn_metadata", + capture_metadata, + ) + speculator = object.__new__(dflash_speculator.DFlashSpeculator) + object.__setattr__(speculator, "draft_attn_layer_names", ["draft"]) + object.__setattr__(speculator, "draft_cp_size", 1) + object.__setattr__(speculator, "block_tables", SimpleNamespace(cp_size=4)) + object.__setattr__( + speculator, + "input_buffers", + SimpleNamespace( + seq_lens=seq_lens, + dcp_local_seq_lens=torch.empty_like(seq_lens), + ), + ) + object.__setattr__(speculator, "num_query_per_req", 4) + + result = speculator._build_draft_attn_metadata( + num_reqs=2, + num_reqs_padded=2, + num_tokens_padded=8, + seq_lens_cpu_upper_bound=seq_lens, + step=0, + ) + + assert result == {} + assert len(captured) == 1 + assert captured[0] is seq_lens diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index 4e67c79494e8..f65c2cc83eed 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -89,6 +89,32 @@ def test_replicated_draft_rejects_backend_without_local_dcp(monkeypatch): cp_utils.check_attention_cp_compatibility(config) +def test_replicated_draft_rejects_missing_backend(monkeypatch): + layer = SimpleNamespace( + impl=SimpleNamespace( + supports_mtp_with_cp_non_trivial_interleave_size=False, + need_to_return_lse_for_decode=False, + ), + get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + ) + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda *_args, **_kwargs: {"draft": layer}, + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ), + speculative_config=SimpleNamespace(method="dflash"), + ) + + with pytest.raises(AssertionError, match="replicated DCP"): + cp_utils.check_attention_cp_compatibility(config) + + def test_attention_cache_spec_errors_are_not_swallowed(monkeypatch): layer_impl = SimpleNamespace( supports_mtp_with_cp_non_trivial_interleave_size=False, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 4ed3e26abd93..8a730eb815f7 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1543,6 +1543,18 @@ 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 """ + dcp_replication_modes = { + spec.dcp_replicated + for spec in kv_cache_spec.values() + if isinstance(spec, AttentionSpec) + } + if len(dcp_replication_modes) > 1: + raise ValueError( + "Hybrid KV cache manager cannot be disabled when attention layers " + "mix DCP-replicated and DCP-sharded KV cache specs. Remove " + "`--disable-hybrid-kv-cache-manager` or disable DCP." + ) + if is_kv_cache_spec_uniform( kv_cache_spec ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec): diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index a078f3324fb5..eeabb8a9cc3b 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -41,7 +41,11 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: if get_spec is not None: spec = get_spec(vllm_config) if getattr(spec, "dcp_replicated", False): - assert backend is None or backend.supports_dcp_replicated, ( + assert backend is not None, ( + "Attention with replicated DCP requires an attention " + "backend that advertises local-DCP support." + ) + assert backend.supports_dcp_replicated, ( "Attention with replicated DCP requires backend support, " f"but {backend.get_name()} does not provide it." ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 27e0754bfc7f..5d04b0525338 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -318,16 +318,19 @@ def _build_draft_attn_metadata( if not self.draft_attn_layer_names: return None assert num_query_per_req is None # Omitted for DFlash, read from self instead - if dcp_local_seq_lens is None and self.draft_cp_size > 1: - prepare_dcp_local_seq_lens( - self.input_buffers.dcp_local_seq_lens, - self.input_buffers.seq_lens, - num_reqs, - self.draft_cp_size, - self.draft_cp_rank, - self.draft_cp_interleave, - ) - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens + if dcp_local_seq_lens is None: + if self.draft_cp_size > 1: + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.draft_cp_size, + self.draft_cp_rank, + self.draft_cp_interleave, + ) + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens + elif getattr(self.block_tables, "cp_size", 1) > 1: + dcp_local_seq_lens = self.input_buffers.seq_lens return super()._build_draft_attn_metadata( num_reqs, num_reqs_padded,