[Bugfix][Spec Decode] Avoid fastsafetensors deadlock for PP draft models (#54416)
Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Canlin Guo committed
Sep 12, 2026 at 14:22 UTC
dca96bf97b350204937a72c4d386b79b425f3989
6 files changed
+42
-2
tests/v1/spec_decode/test_draft_attention_backend_override.py
+8
-1
@@ -6,11 +6,13 @@
6
assert on the config ``load_eagle_model`` hands to ``get_model``.
7
"""
8
9
-from dataclasses import dataclass
9
+from dataclasses import dataclass, field
10
+from types import SimpleNamespace
11
from unittest.mock import patch
12
13
import pytest
14
15
+from vllm.config import LoadConfig
16
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
17
18
@@ -43,6 +45,7 @@ class _VllmConfig:
45
kernel_config: _KernelConfig
46
cache_config: _CacheConfig
47
speculative_config: _SpeculativeConfig
48
+ load_config: LoadConfig = field(default_factory=LoadConfig)
49
50
51
def _config(target_backend: str, draft_backend: str | None) -> _VllmConfig:
@@ -65,6 +68,10 @@ def _capture_draft_config(cfg):
68
69
with (
70
patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model),
71
+ patch(
72
+ "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group",
73
+ return_value=SimpleNamespace(world_size=1),
74
+ ),
75
pytest.raises(_Captured) as exc,
76
):
77
load_eagle_model(object(), cfg)
tests/v1/spec_decode/test_draft_moe_backend_override.py
+8
-1
@@ -8,11 +8,13 @@ target is typically unquantized, and quantized-only backends reject it, so the
8
server fails to start rather than falling back.
9
"""
10
11
-from dataclasses import dataclass
11
+from dataclasses import dataclass, field
12
+from types import SimpleNamespace
13
from unittest.mock import patch
14
15
import pytest
16
17
+from vllm.config import LoadConfig
18
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
19
20
@@ -39,6 +41,7 @@ class _VllmConfig:
41
kernel_config: _KernelConfig
42
cache_config: _CacheConfig
43
speculative_config: _SpeculativeConfig
44
+ load_config: LoadConfig = field(default_factory=LoadConfig)
45
46
47
def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig:
@@ -62,6 +65,10 @@ def _capture_draft_config(cfg):
65
66
with (
67
patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model),
68
+ patch(
69
+ "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group",
70
+ return_value=SimpleNamespace(world_size=1),
71
+ ),
72
pytest.raises(_Captured) as exc,
73
):
74
load_eagle_model(object(), cfg)
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+2
@@ -9,6 +9,7 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
9
get_target_lm_head,
10
maybe_share_target_embed,
11
)
12
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
13
14
15
def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
@@ -37,6 +38,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
38
if speculative_config.kv_cache_dtype is not None
39
else vllm_config.cache_config
40
),
41
+ load_config=get_pp_safe_draft_load_config(vllm_config.load_config),
42
)
43
with set_model_tag("dflash_head"):
44
dflash_model = get_model(
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+2
@@ -6,6 +6,7 @@ import torch.nn as nn
6
from vllm.config import ModelConfig, 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
10
11
logger = init_logger(__name__)
12
@@ -73,6 +74,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
74
if speculative_config.kv_cache_dtype is not None
75
else vllm_config.cache_config
76
),
77
+ load_config=get_pp_safe_draft_load_config(vllm_config.load_config),
78
)
79
# VllmConfig post-init restores the target's quant config because the target
80
# config is retained for DSpark's target-layer metadata, so we must override it.
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+4
@@ -8,6 +8,7 @@ from vllm.distributed.parallel_state import get_pp_group
8
from vllm.lora.layers.base import BaseLayerWithLoRA
9
from vllm.model_executor.model_loader import get_model
10
from vllm.model_executor.models.utils import PPMissingLayer
11
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
12
13
14
def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
@@ -103,6 +104,9 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
104
backend=speculative_config.attention_backend,
105
),
106
)
107
+ draft_load_config = get_pp_safe_draft_load_config(vllm_config.load_config)
108
+ if draft_load_config is not vllm_config.load_config:
109
+ vllm_config = replace(vllm_config, load_config=draft_load_config)
110
with set_model_tag("eagle_head"):
111
eagle_model = get_model(
112
vllm_config=vllm_config, model_config=draft_model_config
vllm/v1/worker/gpu/spec_decode/utils.py
+18
@@ -3,10 +3,28 @@
3
import numpy as np
4
import torch
5
6
+from vllm.config import LoadConfig, replace
7
+from vllm.distributed import get_pp_group
8
+from vllm.logger import init_logger
9
from vllm.v1.outputs import DraftTokenIds
10
from vllm.v1.worker.gpu.async_utils import async_copy_to_np
11
from vllm.v1.worker.gpu.input_batch import InputBatch
12
13
+logger = init_logger(__name__)
14
+
15
+
16
+def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig:
17
+ """Avoid collectives that include PP ranks without a draft model."""
18
+ if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors":
19
+ logger.warning_once(
20
+ "fastsafetensors cannot load a draft model instantiated on only "
21
+ "one pipeline stage; falling back to the standard safetensors "
22
+ "loader for the draft model. The target model still uses "
23
+ "fastsafetensors."
24
+ )
25
+ return replace(load_config, load_format="auto")
26
+ return load_config
27
+
28
29
class DraftTokensHandler:
30
def __init__(self, device: torch.device | None = None):