52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config
|
|
|
|
|
|
@dataclass
|
|
class _FakeEPLBConfig:
|
|
num_redundant_experts: int = 0
|
|
|
|
|
|
@dataclass
|
|
class _FakeParallelConfig:
|
|
pipeline_parallel_size: int = 2
|
|
tensor_parallel_size: int = 8
|
|
enable_eplb: bool = True
|
|
eplb_config: _FakeEPLBConfig = field(
|
|
default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32)
|
|
)
|
|
enable_elastic_ep: bool = True
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.enable_eplb and self.eplb_config.num_redundant_experts:
|
|
raise ValueError("redundant experts require EPLB")
|
|
if self.enable_elastic_ep and not self.enable_eplb:
|
|
raise ValueError("elastic EP requires EPLB")
|
|
|
|
|
|
def test_dspark_parallel_config_disables_eplb_atomically():
|
|
target_config = _FakeParallelConfig()
|
|
|
|
draft_config = _get_dspark_parallel_config(
|
|
target_config,
|
|
tensor_parallel_size=4,
|
|
)
|
|
|
|
assert target_config.pipeline_parallel_size == 2
|
|
assert target_config.tensor_parallel_size == 8
|
|
assert target_config.enable_eplb
|
|
assert target_config.eplb_config.num_redundant_experts == 32
|
|
assert target_config.enable_elastic_ep
|
|
|
|
assert draft_config is not target_config
|
|
assert draft_config.pipeline_parallel_size == 1
|
|
assert draft_config.tensor_parallel_size == 4
|
|
assert not draft_config.enable_eplb
|
|
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
|