[Bugfix][Spec Decode][MoE] Avoid uninitialized EPLB state in DeepSeek V4.1 DSpark drafter (#56387)
Signed-off-by: zack <51604064+luoyuctl@users.noreply.github.com>
zack committed
Sep 15, 2026 at 22:35 UTC
0136df94b0d75732be6c66f620f51289c888e41d
2 files changed
+77
-6
tests/v1/worker/test_dspark_utils.py
new
+51
@@ -0,0 +1,51 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass, field
5
+
6
+from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config
7
+
8
+
9
+@dataclass
10
+class _FakeEPLBConfig:
11
+ num_redundant_experts: int = 0
12
+
13
+
14
+@dataclass
15
+class _FakeParallelConfig:
16
+ pipeline_parallel_size: int = 2
17
+ tensor_parallel_size: int = 8
18
+ enable_eplb: bool = True
19
+ eplb_config: _FakeEPLBConfig = field(
20
+ default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32)
21
+ )
22
+ enable_elastic_ep: bool = True
23
+
24
+ def __post_init__(self) -> None:
25
+ if not self.enable_eplb and self.eplb_config.num_redundant_experts:
26
+ raise ValueError("redundant experts require EPLB")
27
+ if self.enable_elastic_ep and not self.enable_eplb:
28
+ raise ValueError("elastic EP requires EPLB")
29
+
30
+
31
+def test_dspark_parallel_config_disables_eplb_atomically():
32
+ target_config = _FakeParallelConfig()
33
+
34
+ draft_config = _get_dspark_parallel_config(
35
+ target_config,
36
+ tensor_parallel_size=4,
37
+ )
38
+
39
+ assert target_config.pipeline_parallel_size == 2
40
+ assert target_config.tensor_parallel_size == 8
41
+ assert target_config.enable_eplb
42
+ assert target_config.eplb_config.num_redundant_experts == 32
43
+ assert target_config.enable_elastic_ep
44
+
45
+ assert draft_config is not target_config
46
+ assert draft_config.pipeline_parallel_size == 1
47
+ assert draft_config.tensor_parallel_size == 4
48
+ assert not draft_config.enable_eplb
49
+ assert draft_config.eplb_config.num_redundant_experts == 0
50
+ assert not draft_config.enable_elastic_ep
51
+ assert draft_config.eplb_config is not target_config.eplb_config
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+26
-6
@@ -3,7 +3,7 @@
3
4
import torch.nn as nn
5
6
-from vllm.config import ModelConfig, VllmConfig, replace
6
+from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace
7
from vllm.logger import init_logger
8
from vllm.v1.attention.backends.registry import AttentionBackendEnum
9
from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
@@ -31,6 +31,29 @@ def _resolve_dspark_attention_backend(
31
return None
32
33
34
+def _get_dspark_parallel_config(
35
+ parallel_config: ParallelConfig,
36
+ tensor_parallel_size: int,
37
+) -> ParallelConfig:
38
+ if parallel_config.enable_eplb:
39
+ logger.warning_once(
40
+ "EPLB is disabled for the DSpark draft model. EPLB remains enabled "
41
+ "for the target model."
42
+ )
43
+
44
+ return replace(
45
+ parallel_config,
46
+ pipeline_parallel_size=1,
47
+ tensor_parallel_size=tensor_parallel_size,
48
+ enable_eplb=False,
49
+ eplb_config=replace(
50
+ parallel_config.eplb_config,
51
+ num_redundant_experts=0,
52
+ ),
53
+ enable_elastic_ep=False,
54
+ )
55
+
56
+
57
def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
58
speculative_config = vllm_config.speculative_config
59
assert speculative_config is not None
@@ -54,12 +77,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
77
78
draft_vllm_config = replace(
79
vllm_config,
57
- parallel_config=replace(
80
+ parallel_config=_get_dspark_parallel_config(
81
vllm_config.parallel_config,
59
- pipeline_parallel_size=1,
60
- tensor_parallel_size=(
61
- speculative_config.draft_parallel_config.tensor_parallel_size
62
- ),
82
+ speculative_config.draft_parallel_config.tensor_parallel_size,
83
),
84
attention_config=replace(
85
vllm_config.attention_config,