diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 2f8b80a7c1..517184c0f1 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -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, diff --git a/tests/v1/attention/test_group_head_counts.py b/tests/v1/attention/test_group_head_counts.py index 54fe3cfdbc..40d9e2d00f 100644 --- a/tests/v1/attention/test_group_head_counts.py +++ b/tests/v1/attention/test_group_head_counts.py @@ -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 ( diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 24be16e4a5..c2c86053f0 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -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. diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 0a34a6131d..99d3235006 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -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) diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py index 3a474ac203..d2c2ac87b7 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -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.""" diff --git a/tests/v1/kv_connector/unit/test_decode_bench_connector.py b/tests/v1/kv_connector/unit/test_decode_bench_connector.py index 693c10d40a..6f518ab53a 100644 --- a/tests/v1/kv_connector/unit/test_decode_bench_connector.py +++ b/tests/v1/kv_connector/unit/test_decode_bench_connector.py @@ -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, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index ed07718de6..3817a15e3e 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -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. diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 16d6d8e516..1d004c3667 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -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] diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index b542388952..2c4b951be7 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -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] diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index b38ae4d063..d477ace479 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -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( diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py index 94655f192c..1014051d64 100644 --- a/tests/v1/worker/test_dspark_utils.py +++ b/tests/v1/worker/test_dspark_utils.py @@ -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) diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 4fc7d4d2e5..36ba9a4b16 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -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, ), ) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 5c245d1c4d..5c2231a1f2 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -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(): diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index 0c4b5aff68..4612d49d92 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -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, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py index a9255f7cda..cbb7339c91 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py @@ -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 ) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 20118649e1..88bb3d7c5a 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -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 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index b5cd6e4b32..5d661634d7 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -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 diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 129b3e39c3..6035678394 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -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 diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index f40c75b474..b3d40b1c7f 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -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, diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index e918774088..c97483912d 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -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: diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 8061999409..d69b2f5fb3 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -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; diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index ee0957839e..3f4f4e7bb1 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -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))) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 0dd9734847..7352f71cb6 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -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 = ( diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 499b50a3a2..5791704b2b 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -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) diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 96bfb35a67..344daa36b8 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -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): diff --git a/vllm/v1/spec_decode/draft_model.py b/vllm/v1/spec_decode/draft_model.py index 08542f03a1..17be2732e6 100644 --- a/vllm/v1/spec_decode/draft_model.py +++ b/vllm/v1/spec_decode/draft_model.py @@ -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, ) diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 79651759aa..bcf56dc0bf 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -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 " diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 096d9b99aa..8cf3ba98dd 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 61ceaeef2e..9f25425ab6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -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, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index aff8f04cf3..de49fcffb9 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -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}." diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index cdcdd5fcba..57110e0b05 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -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, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d1d6d888aa..6f88febb66 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -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, )