| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass, field |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import ParallelConfig |
| 11 | from vllm.v1.kv_cache_interface import FullAttentionSpec |
| 12 | from vllm.v1.worker.gpu import model_runner |
| 13 | from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator |
| 14 | from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config |
| 15 | from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator |
| 16 | |
| 17 | |
| 18 | @dataclass |
| 19 | class _FakeEPLBConfig: |
| 20 | num_redundant_experts: int = 0 |
| 21 | |
| 22 | |
| 23 | @dataclass |
| 24 | class _FakeParallelConfig: |
| 25 | pipeline_parallel_size: int = 2 |
| 26 | tensor_parallel_size: int = 8 |
| 27 | enable_eplb: bool = True |
| 28 | eplb_config: _FakeEPLBConfig = field( |
| 29 | default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32) |
| 30 | ) |
| 31 | enable_elastic_ep: bool = True |
| 32 | |
| 33 | def __post_init__(self) -> None: |
| 34 | if not self.enable_eplb and self.eplb_config.num_redundant_experts: |
| 35 | raise ValueError("redundant experts require EPLB") |
| 36 | if self.enable_elastic_ep and not self.enable_eplb: |
| 37 | raise ValueError("elastic EP requires EPLB") |
| 38 | |
| 39 | |
| 40 | def test_dspark_parallel_config_disables_eplb_atomically(): |
| 41 | target_config = _FakeParallelConfig() |
| 42 | |
| 43 | draft_config = _get_dspark_parallel_config( |
| 44 | target_config, |
| 45 | tensor_parallel_size=4, |
| 46 | ) |
| 47 | |
| 48 | assert target_config.pipeline_parallel_size == 2 |
| 49 | assert target_config.tensor_parallel_size == 8 |
| 50 | assert target_config.enable_eplb |
| 51 | assert target_config.eplb_config.num_redundant_experts == 32 |
| 52 | assert target_config.enable_elastic_ep |
| 53 | |
| 54 | assert draft_config is not target_config |
| 55 | assert draft_config.pipeline_parallel_size == 1 |
| 56 | assert draft_config.tensor_parallel_size == 4 |
| 57 | assert not draft_config.enable_eplb |
| 58 | assert draft_config.eplb_config.num_redundant_experts == 0 |
| 59 | assert not draft_config.enable_elastic_ep |
| 60 | assert draft_config.eplb_config is not target_config.eplb_config |
| 61 | |
| 62 | |
| 63 | @pytest.mark.parametrize("pcp_size", [1, 4]) |
| 64 | @pytest.mark.parametrize("dcp_size", [1, 4]) |
| 65 | @pytest.mark.parametrize("use_mla", [False, True]) |
| 66 | def test_draft_context_parallelism_without_changing_target( |
| 67 | monkeypatch, pcp_size, dcp_size, use_mla |
| 68 | ): |
| 69 | target_parallel = ParallelConfig( |
| 70 | tensor_parallel_size=4, |
| 71 | prefill_context_parallel_size=pcp_size, |
| 72 | decode_context_parallel_size=dcp_size, |
| 73 | cp_kv_cache_interleave_size=16, |
| 74 | distributed_executor_backend="mp", |
| 75 | ) |
| 76 | target_config = SimpleNamespace( |
| 77 | parallel_config=target_parallel, |
| 78 | speculative_config=SimpleNamespace( |
| 79 | draft_model_config=SimpleNamespace(use_mla=use_mla) |
| 80 | ), |
| 81 | ) |
| 82 | |
| 83 | class CapturedConfig(Exception): |
| 84 | pass |
| 85 | |
| 86 | def capture_init(self, config, device): |
| 87 | raise CapturedConfig(config) |
| 88 | |
| 89 | monkeypatch.setattr(DraftModelSpeculator, "__init__", capture_init) |
| 90 | with pytest.raises(CapturedConfig) as captured: |
| 91 | DFlashSpeculator(target_config, device=None) |
| 92 | draft_parallel = captured.value.args[0].parallel_config |
| 93 | assert draft_parallel.tensor_parallel_size == 4 |
| 94 | assert draft_parallel.prefill_context_parallel_size == 1 |
| 95 | assert draft_parallel.decode_context_parallel_size == (dcp_size if use_mla else 1) |
| 96 | assert draft_parallel.cp_kv_cache_interleave_size == 16 |
| 97 | assert target_parallel.prefill_context_parallel_size == pcp_size |
| 98 | assert target_parallel.decode_context_parallel_size == dcp_size |
| 99 | assert target_parallel.cp_kv_cache_interleave_size == 16 |
| 100 | |
| 101 | |
| 102 | @pytest.mark.parametrize("configured_dcp", [1, 4]) |
| 103 | def test_attention_uses_draft_dcp_setting_inside_target_process_group( |
| 104 | monkeypatch, configured_dcp |
| 105 | ): |
| 106 | import vllm.config as config_module |
| 107 | from vllm.distributed import parallel_state |
| 108 | from vllm.v1.attention.backend import AttentionImplBase |
| 109 | |
| 110 | config = SimpleNamespace( |
| 111 | parallel_config=SimpleNamespace(decode_context_parallel_size=configured_dcp) |
| 112 | ) |
| 113 | monkeypatch.setattr( |
| 114 | config_module, "get_current_vllm_config_or_none", lambda: config |
| 115 | ) |
| 116 | |
| 117 | monkeypatch.setattr( |
| 118 | parallel_state, "_DCP", SimpleNamespace(world_size=4, rank_in_group=2) |
| 119 | ) |
| 120 | impl = AttentionImplBase() |
| 121 | assert impl.dcp_world_size == configured_dcp |
| 122 | assert impl.dcp_rank == (0 if configured_dcp == 1 else 2) |
| 123 | |
| 124 | |
| 125 | @pytest.mark.parametrize("draft_dcp_size", [1, 4]) |
| 126 | def test_runner_marks_only_replicated_draft_caches(monkeypatch, draft_dcp_size): |
| 127 | spec = FullAttentionSpec( |
| 128 | block_size=16, num_kv_heads=1, head_size=64, dtype=torch.bfloat16 |
| 129 | ) |
| 130 | monkeypatch.setattr( |
| 131 | model_runner, "get_kv_cache_spec", lambda _: {"target": spec, "draft": spec} |
| 132 | ) |
| 133 | runner = object.__new__(model_runner.GPUModelRunner) |
| 134 | runner.vllm_config = None |
| 135 | runner.dcp_size = 4 |
| 136 | runner.speculator = object.__new__(DFlashSpeculator) |
| 137 | runner.speculator.dcp_size = draft_dcp_size |
| 138 | runner.speculator.draft_attn_layer_names = {"draft"} |
| 139 | |
| 140 | specs = runner.get_kv_cache_spec() |
| 141 | |
| 142 | assert specs["target"].dcp_sharded |
| 143 | assert specs["draft"].dcp_sharded == (draft_dcp_size == 4) |