[Bugfix][Spec Decode] Honour the draft's attention_backend on Model Runner V2 (#54826)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Stefano Castagnetta committed
Sep 4, 2026 at 21:23 UTC
25240513856e9aca9b07e9381f2a4513b6954bd8
2 files changed
+100
tests/v1/spec_decode/test_draft_attention_backend_override.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The draft must honour ``attention_backend`` from --speculative-config.
4
+
5
+``init_attn_backend`` reads the backend off each constructed layer, so these
6
+assert on the config ``load_eagle_model`` hands to ``get_model``.
7
+"""
8
+
9
+from dataclasses import dataclass
10
+from unittest.mock import patch
11
+
12
+import pytest
13
+
14
+from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
15
+
16
+
17
+@dataclass
18
+class _AttentionConfig:
19
+ backend: str | None = None
20
+
21
+
22
+@dataclass
23
+class _KernelConfig:
24
+ moe_backend: str | None = None
25
+
26
+
27
+@dataclass
28
+class _CacheConfig:
29
+ cache_dtype: str = "auto"
30
+
31
+
32
+@dataclass
33
+class _SpeculativeConfig:
34
+ attention_backend: str | None = None
35
+ moe_backend: str | None = None
36
+ kv_cache_dtype: str | None = None
37
+ draft_model_config: object = None
38
+
39
+
40
+@dataclass
41
+class _VllmConfig:
42
+ attention_config: _AttentionConfig
43
+ kernel_config: _KernelConfig
44
+ cache_config: _CacheConfig
45
+ speculative_config: _SpeculativeConfig
46
+
47
+
48
+def _config(target_backend: str, draft_backend: str | None) -> _VllmConfig:
49
+ return _VllmConfig(
50
+ attention_config=_AttentionConfig(backend=target_backend),
51
+ kernel_config=_KernelConfig(),
52
+ cache_config=_CacheConfig(),
53
+ speculative_config=_SpeculativeConfig(attention_backend=draft_backend),
54
+ )
55
+
56
+
57
+class _Captured(Exception):
58
+ def __init__(self, vllm_config):
59
+ self.vllm_config = vllm_config
60
+
61
+
62
+def _capture_draft_config(cfg):
63
+ def _fake_get_model(*, vllm_config, model_config):
64
+ raise _Captured(vllm_config)
65
+
66
+ with (
67
+ patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model),
68
+ pytest.raises(_Captured) as exc,
69
+ ):
70
+ load_eagle_model(object(), cfg)
71
+ return exc.value.vllm_config
72
+
73
+
74
+def test_draft_attention_backend_overrides_the_target():
75
+ used = _capture_draft_config(_config("FLASHINFER", "TRITON_ATTN"))
76
+ assert used.attention_config.backend == "TRITON_ATTN"
77
+
78
+
79
+def test_unset_leaves_the_target_backend_in_place():
80
+ """Clearing it lets the draft autoselect a KV layout the target lacks."""
81
+ cfg = _config("FLEX_ATTENTION", None)
82
+ used = _capture_draft_config(cfg)
83
+ assert used is cfg
84
+ assert used.attention_config.backend == "FLEX_ATTENTION"
85
+
86
+
87
+def test_override_does_not_mutate_the_target_config():
88
+ cfg = _config("FLASHINFER", "TRITON_ATTN")
89
+ _capture_draft_config(cfg)
90
+ assert cfg.attention_config.backend == "FLASHINFER"
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+10
@@ -47,6 +47,16 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
47
cache_dtype=speculative_config.kv_cache_dtype,
48
),
49
)
50
+ if speculative_config.attention_backend is not None:
51
+ # Before get_model(): the backend is read off the constructed layers.
52
+ # Only when set, so the draft keeps a KV cache layout the target shares.
53
+ vllm_config = replace(
54
+ vllm_config,
55
+ attention_config=replace(
56
+ vllm_config.attention_config,
57
+ backend=speculative_config.attention_backend,
58
+ ),
59
+ )
60
with set_model_tag("eagle_head"):
61
eagle_model = get_model(
62
vllm_config=vllm_config, model_config=draft_model_config