[Bugfix][Spec Decode] Scope DSpark backend inheritance to DeepSeek V4 (#52809)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Misha Goin committed
Aug 20, 2026 at 21:42 UTC
91a893de64722019ea2faf852e06cabe143b3490
1 file changed
+34
-10
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+34
-10
@@ -3,13 +3,32 @@
3
4
import torch.nn as nn
5
6
-from vllm.config import VllmConfig, replace
6
+from vllm.config import ModelConfig, VllmConfig, replace
7
from vllm.distributed.parallel_state import get_pp_group
8
-from vllm.model_executor.model_loader import get_model
9
-from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
10
- _should_share,
11
- get_target_lm_head,
12
-)
8
+from vllm.logger import init_logger
9
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
10
+
11
+logger = init_logger(__name__)
12
+
13
+
14
+def _resolve_dspark_attention_backend(
15
+ draft_model_config: ModelConfig,
16
+ draft_backend: AttentionBackendEnum | None,
17
+ target_backend: AttentionBackendEnum | None,
18
+) -> AttentionBackendEnum | None:
19
+ if draft_backend is not None:
20
+ return draft_backend
21
+ # DeepSeek-V4 draft layers share the target's KV-cache layout. Other
22
+ # DSpark architectures may use a different attention kind.
23
+ if draft_model_config.hf_config.model_type == "deepseek_v4":
24
+ if target_backend is not None:
25
+ logger.info_once(
26
+ "Using the target model's %s attention backend for the "
27
+ "DeepSeek-V4 DSpark drafter.",
28
+ target_backend.name,
29
+ )
30
+ return target_backend
31
+ return None
32
33
34
def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
@@ -18,13 +37,18 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
37
draft_model_config = speculative_config.draft_model_config
38
39
from vllm.compilation.backends import set_model_tag
40
+ from vllm.model_executor.model_loader import get_model
41
from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
42
from vllm.model_executor.models.utils import get_draft_quant_config
43
+ from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
44
+ _should_share,
45
+ get_target_lm_head,
46
+ )
47
24
- # None re-runs backend auto-selection for the draft, which can pick a
25
- # different attention class than the target; fall back to the target's.
26
- draft_attention_backend = (
27
- speculative_config.attention_backend or vllm_config.attention_config.backend
48
+ draft_attention_backend = _resolve_dspark_attention_backend(
49
+ draft_model_config,
50
+ speculative_config.attention_backend,
51
+ vllm_config.attention_config.backend,
52
)
53
54
draft_vllm_config = replace(