[Bugfix][Spec Decode] Honour the draft's moe_backend on Model Runner V2 (#54788)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Stefano Castagnetta committed
Sep 9, 2026 at 01:50 UTC
9c2d21046bb39f56ee93a5fab4f899714a5579ef
2 files changed
+97
tests/v1/spec_decode/test_draft_moe_backend_override.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The draft model must honour ``moe_backend`` from --speculative-config.
4
+
5
+The draft is loaded with the target's VllmConfig, so without an explicit
6
+override it inherits the target's --moe-backend. An MTP head on a quantized
7
+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
12
+from unittest.mock import patch
13
+
14
+import pytest
15
+
16
+from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
17
+
18
+
19
+@dataclass
20
+class _KernelConfig:
21
+ moe_backend: str | None = None
22
+
23
+
24
+@dataclass
25
+class _CacheConfig:
26
+ cache_dtype: str = "auto"
27
+
28
+
29
+@dataclass
30
+class _SpeculativeConfig:
31
+ attention_backend: str | None = None
32
+ moe_backend: str | None = None
33
+ kv_cache_dtype: str | None = None
34
+ draft_model_config: object = None
35
+
36
+
37
+@dataclass
38
+class _VllmConfig:
39
+ kernel_config: _KernelConfig
40
+ cache_config: _CacheConfig
41
+ speculative_config: _SpeculativeConfig
42
+
43
+
44
+def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig:
45
+ return _VllmConfig(
46
+ kernel_config=_KernelConfig(moe_backend=target_moe),
47
+ cache_config=_CacheConfig(),
48
+ speculative_config=_SpeculativeConfig(moe_backend=draft_moe),
49
+ )
50
+
51
+
52
+class _Captured(Exception):
53
+ """Stops load_eagle_model once we hold the config it would have used."""
54
+
55
+ def __init__(self, vllm_config):
56
+ self.vllm_config = vllm_config
57
+
58
+
59
+def _capture_draft_config(cfg):
60
+ def _fake_get_model(*, vllm_config, model_config):
61
+ raise _Captured(vllm_config)
62
+
63
+ with (
64
+ patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model),
65
+ pytest.raises(_Captured) as exc,
66
+ ):
67
+ load_eagle_model(object(), cfg)
68
+ return exc.value.vllm_config
69
+
70
+
71
+def test_draft_moe_backend_overrides_the_target():
72
+ """A draft moe_backend must reach the draft's kernel config."""
73
+ used = _capture_draft_config(_config("flashinfer_b12x", "flashinfer_cutlass"))
74
+ assert used.kernel_config.moe_backend == "flashinfer_cutlass"
75
+
76
+
77
+def test_draft_inherits_target_when_no_override():
78
+ """No override means the draft still inherits the target, as before."""
79
+ used = _capture_draft_config(_config("flashinfer_b12x", None))
80
+ assert used.kernel_config.moe_backend == "flashinfer_b12x"
81
+
82
+
83
+def test_override_does_not_mutate_the_target_config():
84
+ """The target must keep its own backend after the draft is built."""
85
+ cfg = _config("flashinfer_b12x", "flashinfer_cutlass")
86
+ _capture_draft_config(cfg)
87
+ assert cfg.kernel_config.moe_backend == "flashinfer_b12x"
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+10
@@ -75,6 +75,16 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
75
speculative_config = vllm_config.speculative_config
76
assert speculative_config is not None
77
draft_model_config = speculative_config.draft_model_config
78
+ if speculative_config.moe_backend is not None:
79
+ # Otherwise the draft inherits the target's --moe-backend, which
80
+ # fails when the draft is unquantized and that backend is not.
81
+ vllm_config = replace(
82
+ vllm_config,
83
+ kernel_config=replace(
84
+ vllm_config.kernel_config,
85
+ moe_backend=speculative_config.moe_backend,
86
+ ),
87
+ )
88
if speculative_config.kv_cache_dtype is not None:
89
vllm_config = replace(
90
vllm_config,