A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 00:45 UTC
[Feature] Support per-draft-model MoE backend via `--speculative-config` (#37880)
Signed-off-by: Andrii Skliar <askliar@nvidia.com> Signed-off-by: [Andrii Skliar] <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>
Andrii Skliar committed
Mar 25, 2026 at 15:31 UTC
cd7643015e583c1e78d437118a6ce8282cb85663
5 files changed
+140
-40
tests/v1/e2e/spec_decode/test_spec_decode.py
+98
-8
@@ -20,12 +20,11 @@ from vllm import LLM, SamplingParams
20
from vllm.assets.base import VLLM_S3_BUCKET_URL
21
from vllm.assets.image import VLM_IMAGES_DIR
22
from vllm.benchmarks.datasets import InstructCoderDataset
23
-from vllm.config import VllmConfig
23
+from vllm.config import VllmConfig, replace
24
from vllm.distributed import cleanup_dist_env_and_memory
25
from vllm.engine.arg_utils import EngineArgs
26
from vllm.platforms import current_platform
27
from vllm.v1.metrics.reader import Metric
28
-from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model
28
29
MTP_SIMILARITY_RATE = 0.8
30
@@ -919,13 +918,104 @@ def test_draft_model_engine_args_tensor_parallelism():
918
"draft_tensor_parallel_size": 1, # <<< valid arg name
919
},
920
)
922
- tgt_vllm_config: VllmConfig = engine_args.create_engine_config()
923
- assert tgt_vllm_config.parallel_config.tensor_parallel_size == 2
924
- assert tgt_vllm_config.quant_config.get_name() == "fp8"
921
+ target_config: VllmConfig = engine_args.create_engine_config()
922
+ assert target_config.parallel_config.tensor_parallel_size == 2
923
+ assert target_config.quant_config.get_name() == "fp8"
924
+
925
+ speculative_config = target_config.speculative_config
926
+ draft_config: VllmConfig = replace(
927
+ target_config,
928
+ quant_config=None,
929
+ parallel_config=replace(
930
+ speculative_config.draft_parallel_config,
931
+ rank=target_config.parallel_config.rank,
932
+ ),
933
+ model_config=speculative_config.draft_model_config,
934
+ )
935
+ assert draft_config.parallel_config.tensor_parallel_size == 1
936
+ assert draft_config.quant_config is None
937
+
938
+
939
+def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
940
+ """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic
941
+ so we can test it without instantiating a full proposer."""
942
+ spec_cfg = vllm_config.speculative_config
943
+ if spec_cfg.moe_backend is not None:
944
+ return replace(
945
+ vllm_config,
946
+ kernel_config=replace(
947
+ vllm_config.kernel_config,
948
+ moe_backend=spec_cfg.moe_backend,
949
+ ),
950
+ )
951
+ return vllm_config
952
+
953
+
954
+def test_draft_model_moe_backend_override():
955
+ """When moe_backend is set in speculative_config, the draft VllmConfig
956
+ should use it while the target keeps its own setting."""
957
+ engine_args = EngineArgs(
958
+ model="Qwen/Qwen3-1.7B",
959
+ tensor_parallel_size=1,
960
+ moe_backend="flashinfer_trtllm",
961
+ speculative_config={
962
+ "model": "Qwen/Qwen3-0.6B",
963
+ "method": "draft_model",
964
+ "num_speculative_tokens": 3,
965
+ "moe_backend": "triton",
966
+ },
967
+ )
968
+ tgt_config: VllmConfig = engine_args.create_engine_config()
969
+ assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
970
+ assert tgt_config.speculative_config.moe_backend == "triton"
971
+
972
+ draft_config = _apply_draft_moe_backend(tgt_config)
973
+ assert draft_config.kernel_config.moe_backend == "triton"
974
+ # Target config must be unaffected.
975
+ assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
976
+
977
+
978
+def test_draft_model_moe_backend_inherits_target():
979
+ """When moe_backend is not set in speculative_config, the draft should
980
+ inherit the target's moe_backend."""
981
+ engine_args = EngineArgs(
982
+ model="Qwen/Qwen3-1.7B",
983
+ tensor_parallel_size=1,
984
+ moe_backend="flashinfer_cutlass",
985
+ speculative_config={
986
+ "model": "Qwen/Qwen3-0.6B",
987
+ "method": "draft_model",
988
+ "num_speculative_tokens": 3,
989
+ },
990
+ )
991
+ tgt_config: VllmConfig = engine_args.create_engine_config()
992
+ assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass"
993
+ assert tgt_config.speculative_config.moe_backend is None
994
+
995
+ draft_config = _apply_draft_moe_backend(tgt_config)
996
+ assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass"
997
+ assert draft_config is tgt_config
998
+
999
+
1000
+def test_draft_model_moe_backend_default_auto():
1001
+ """When neither target nor draft set moe_backend explicitly, both should
1002
+ default to 'auto'."""
1003
+ engine_args = EngineArgs(
1004
+ model="Qwen/Qwen3-1.7B",
1005
+ tensor_parallel_size=1,
1006
+ speculative_config={
1007
+ "model": "Qwen/Qwen3-0.6B",
1008
+ "method": "draft_model",
1009
+ "num_speculative_tokens": 3,
1010
+ },
1011
+ )
1012
+ tgt_config: VllmConfig = engine_args.create_engine_config()
1013
+ assert tgt_config.kernel_config.moe_backend == "auto"
1014
+ assert tgt_config.speculative_config.moe_backend is None
1015
926
- draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(tgt_vllm_config)
927
- assert draft_vllm_config.parallel_config.tensor_parallel_size == 1
928
- assert draft_vllm_config.quant_config is None
1016
+ draft_config = _apply_draft_moe_backend(tgt_config)
1017
+ assert draft_config.kernel_config.moe_backend == "auto"
1018
+ assert draft_config is tgt_config
1019
1020
1021
def test_draft_model_engine_args_rejects_invalid_tp_argname():
vllm/config/speculative.py
+6
@@ -9,6 +9,7 @@ from pydantic import Field, SkipValidation, model_validator
9
from typing_extensions import Self
10
11
from vllm.config import LoadConfig
12
+from vllm.config.kernel import MoEBackend
13
from vllm.config.model import ModelConfig
14
from vllm.config.parallel import ParallelConfig
15
from vllm.config.utils import config
@@ -93,6 +94,11 @@ class SpeculativeConfig:
94
"""Quantization method that was used to quantize the draft model weights.
95
If `None`, we assume the model weights are not quantized. Note that it only
96
takes effect when using the draft model-based speculative method."""
97
+ moe_backend: MoEBackend | None = None
98
+ """MoE backend to use for the draft model. When `None`, the draft model
99
+ inherits the target model's `--moe-backend` setting. Useful when the
100
+ drafter and generator require different MoE kernels (e.g. quantized
101
+ generator with unquantized drafter)."""
102
max_model_len: int | None = Field(default=None, ge=1)
103
"""The maximum model length of the draft model. Used when testing the
104
ability to skip speculation for some sequences."""
vllm/v1/spec_decode/draft_model.py
+18
-5
@@ -6,10 +6,10 @@ import torch.nn as nn
6
from typing_extensions import override
7
8
from vllm.config import VllmConfig
9
+from vllm.config.utils import replace
10
from vllm.logger import init_logger
11
from vllm.model_executor.model_loader import get_model
12
from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer
12
-from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model
13
14
logger = init_logger(__name__)
15
@@ -50,16 +50,29 @@ class DraftModelProposer(SpecDecodeBaseProposer):
50
"Please pass 'draft_tensor_parallel_size' in the speculative_config."
51
)
52
53
+ @override
54
+ def _create_draft_vllm_config(self) -> VllmConfig:
55
+ base = super()._create_draft_vllm_config()
56
+ spec = self.speculative_config
57
+
58
+ return replace(
59
+ base,
60
+ quant_config=None,
61
+ parallel_config=replace(
62
+ spec.draft_parallel_config,
63
+ rank=self.vllm_config.parallel_config.rank,
64
+ ),
65
+ model_config=spec.draft_model_config,
66
+ )
67
+
68
@override
69
def _get_model(self) -> nn.Module:
55
- # Draft models may be quantized or on different parallelism,
56
- # so we load them with a modified vllm config
70
from vllm.compilation.backends import set_model_tag
71
59
- temp_vllm_config = create_vllm_config_for_draft_model(self.vllm_config)
72
+ draft_vllm_config = self._create_draft_vllm_config()
73
with set_model_tag("draft_model"):
74
model = get_model(
62
- vllm_config=temp_vllm_config,
75
+ vllm_config=draft_vllm_config,
76
prefix="draft_model",
77
)
78
return model
vllm/v1/spec_decode/eagle.py
+18
-2
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import ast
4
-from dataclasses import replace
4
from importlib.util import find_spec
5
from typing import cast
6
@@ -13,6 +12,7 @@ from vllm.config import (
12
CUDAGraphMode,
13
VllmConfig,
14
get_layers_from_vllm_config,
15
+ replace,
16
)
17
from vllm.distributed.parallel_state import get_pp_group
18
from vllm.forward_context import set_forward_context
@@ -1213,6 +1213,21 @@ class SpecDecodeBaseProposer:
1213
model = model.module
1214
return model.__class__.__name__
1215
1216
+ def _create_draft_vllm_config(self) -> VllmConfig:
1217
+ """Return a VllmConfig with kernel-level overrides for the proposer.
1218
+ Subclasses may override to apply additional config changes.
1219
+ """
1220
+ spec_cfg = self.speculative_config
1221
+ if spec_cfg.moe_backend is not None:
1222
+ return replace(
1223
+ self.vllm_config,
1224
+ kernel_config=replace(
1225
+ self.vllm_config.kernel_config,
1226
+ moe_backend=spec_cfg.moe_backend,
1227
+ ),
1228
+ )
1229
+ return self.vllm_config
1230
+
1231
def _get_model(self) -> nn.Module:
1232
"""
1233
Default method to call get_model(). Can be overridden by subclasses which
@@ -1220,9 +1235,10 @@ class SpecDecodeBaseProposer:
1235
"""
1236
from vllm.compilation.backends import set_model_tag
1237
1238
+ draft_vllm_config = self._create_draft_vllm_config()
1239
with set_model_tag("eagle_head"):
1240
model = get_model(
1225
- vllm_config=self.vllm_config,
1241
+ vllm_config=draft_vllm_config,
1242
model_config=self.speculative_config.draft_model_config,
1243
load_config=self.speculative_config.draft_load_config,
1244
)
vllm/v1/spec_decode/utils.py
-25
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-from vllm.config import VllmConfig, replace
5
from vllm.platforms import current_platform
6
from vllm.triton_utils import tl, triton
7
from vllm.v1.attention.backends.utils import (
@@ -258,30 +257,6 @@ def compute_new_slot_mapping(
257
return new_slot_mapping
258
259
261
-def create_vllm_config_for_draft_model(
262
- target_model_vllm_config: VllmConfig,
263
-) -> VllmConfig:
264
- """The vllm_config is configured for the target model, e.g.
265
- its quant_config and parallel_config. But the draft model is potentially
266
- quantized differently, and has potentially different tensor_parallel_size.
267
- This function creates a new vllm_config configured for the drafter.
268
- The vllm_config is useful when loading the draft model with get_model().
269
- """
270
- old = target_model_vllm_config
271
- assert old.speculative_config is not None, "speculative_config is not set"
272
- old_spec_config = old.speculative_config
273
- new_parallel_config = replace(
274
- old_spec_config.draft_parallel_config, rank=old.parallel_config.rank
275
- )
276
- new: VllmConfig = replace(
277
- old,
278
- quant_config=None,
279
- parallel_config=new_parallel_config,
280
- model_config=old_spec_config.draft_model_config,
281
- )
282
- return new
283
-
284
-
260
def extend_all_queries_by_N(
261
common_attn_metadata: CommonAttentionMetadata,
262
N: int,