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 3, 19:00 UTC
[Feature] Support custom callable proposer backend for speculative decoding (#39487)
Signed-off-by: 524031910363 <hyzhyzsh@sjtu.edu.cn> Signed-off-by: CynicDora <hyzhyzsh@sjtu.edu.cn>
CynicDora committed
May 14, 2026 at 00:53 UTC
256dbcaabf2acc38dad804e8beac7ea0887e72d4
7 files changed
+261
-11
docs/features/speculative_decoding/README.md
+12
@@ -15,6 +15,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
15
- [Multi-Layer Perceptron](mlp.md)
16
- [N-Gram](n_gram.md)
17
- [Suffix Decoding](suffix.md)
18
+- [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)
19
20
## Method Selection at a Glance
21
@@ -30,11 +31,22 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
31
| MLP speculator | Medium to high gain | Medium gain | Good when compatible MLP speculators are available. |
32
| N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. |
33
| Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. |
34
+| Custom Proposer | Varies | Varies | Bring your own proposer class (experimental). |
35
36
For reproducible measurements in your environment, use
37
[`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)
38
or the [benchmark CLI guide](../../benchmarking/cli.md).
39
40
+## Custom Proposer Backend (Experimental)
41
+
42
+You can plug in your own custom proposer class for speculative decoding by setting the method to `custom_class` and providing the full module path to your class.
43
+Your custom class must accept a `VllmConfig` upon instantiation and implement a `propose` method.
44
+
45
+**Example configuration:**
46
+
47
+- `speculative_config.method = "custom_class"`
48
+- `speculative_config.model = "your_module.YourCustomProposerClass"`
49
+
50
## `--speculative-config` schema
51
52
Use `--speculative-config` to pass speculative decoding settings as a JSON
tests/spec_decode/test_custom_proposer.py
new
+121
@@ -0,0 +1,121 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Integration test for custom proposer class in speculative decoding.
5
+
6
+Usage:
7
+ .venv/bin/python test_custom_proposer.py
8
+"""
9
+
10
+import os
11
+
12
+import torch
13
+
14
+from vllm import LLM, SamplingParams
15
+from vllm.config import VllmConfig
16
+
17
+MODEL_ID = "facebook/opt-125m"
18
+NUM_SPEC_TOKENS = 5
19
+
20
+
21
+class DummyDraftProposer:
22
+ """Custom proposer class that repeats the last token of each sequence.
23
+
24
+ This demonstrates the class-based custom proposer interface.
25
+ """
26
+
27
+ def __init__(self, vllm_config: VllmConfig):
28
+ """Initialize the custom proposer.
29
+
30
+ Args:
31
+ vllm_config: vLLM configuration containing model and speculative settings.
32
+ """
33
+ self.num_speculative_tokens = (
34
+ vllm_config.speculative_config.num_speculative_tokens
35
+ )
36
+ self.max_model_len = vllm_config.model_config.max_model_len
37
+ print(
38
+ f"[DummyDraftProposer.__init__] num_speculative_tokens="
39
+ f"{self.num_speculative_tokens}, max_model_len={self.max_model_len}"
40
+ )
41
+
42
+ def propose(
43
+ self,
44
+ sampled_token_ids: list[list[int]],
45
+ num_tokens_no_spec: int,
46
+ token_ids_cpu: torch.Tensor,
47
+ slot_mappings: torch.Tensor | None = None,
48
+ ) -> list[list[int]]:
49
+ """Generate draft tokens by repeating the last token of each sequence.
50
+
51
+ Args:
52
+ sampled_token_ids: Recently sampled token IDs per request.
53
+ num_tokens_no_spec: Number of non-speculative tokens per request.
54
+ token_ids_cpu: Full token IDs tensor on CPU.
55
+ slot_mappings: Slot mapping for KV cache (optional).
56
+
57
+ Returns:
58
+ List of draft token sequences for each request.
59
+ """
60
+ # Cross-process flag to prove this method was executed
61
+ with open("proposer_called.flag", "w") as f:
62
+ f.write("called")
63
+
64
+ batch_size = len(sampled_token_ids)
65
+ last_tokens = [seq[-1] for seq in sampled_token_ids]
66
+ drafts = [[t] * self.num_speculative_tokens for t in last_tokens]
67
+ print(
68
+ f"[DummyDraftProposer.propose] batch_size={batch_size}, "
69
+ f"num_speculative_tokens={self.num_speculative_tokens}, "
70
+ f"drafts_shape={len(drafts)}x{len(drafts[0])}"
71
+ )
72
+ return drafts
73
+
74
+
75
+if __name__ == "__main__":
76
+ print("=" * 60)
77
+ print("Custom Proposer Backend Integration Test")
78
+ print("=" * 60)
79
+
80
+ # Cleanup any leftover flag from previous failed runs
81
+ if os.path.exists("proposer_called.flag"):
82
+ os.remove("proposer_called.flag")
83
+
84
+ llm = LLM(
85
+ model=MODEL_ID,
86
+ speculative_config={
87
+ "model": f"{__name__}.DummyDraftProposer",
88
+ "num_speculative_tokens": NUM_SPEC_TOKENS,
89
+ },
90
+ gpu_memory_utilization=0.4,
91
+ enforce_eager=True,
92
+ )
93
+
94
+ prompts = [
95
+ "Hello, my name is",
96
+ "The future of AI is",
97
+ ]
98
+
99
+ sampling_params = SamplingParams(
100
+ max_tokens=32,
101
+ temperature=0.0,
102
+ )
103
+
104
+ print(f"\nRunning generate with {len(prompts)} prompt(s)...\n")
105
+ outputs = llm.generate(prompts, sampling_params)
106
+
107
+ for output in outputs:
108
+ prompt = output.prompt
109
+ generated = output.outputs[0].text
110
+ print(f"Prompt: {prompt!r}")
111
+ print(f"Generated text: {generated!r}")
112
+ print("-" * 60)
113
+
114
+ # Verify the custom proposer's propose() was actually called across processes
115
+ assert os.path.exists("proposer_called.flag"), (
116
+ "The custom proposer's propose() method was never called!"
117
+ )
118
+ os.remove("proposer_called.flag")
119
+
120
+ print("✓ Custom proposer was actively used during generation!")
121
+ print("Test completed successfully.")
tools/pre_commit/mypy.py
+1
-2
@@ -98,7 +98,6 @@ def mypy(
98
99
100
def main():
101
- ci = sys.argv[1] == "1"
101
python_version = sys.argv[2]
102
file_groups = group_files(sys.argv[3:])
103
@@ -107,7 +106,7 @@ def main():
106
107
returncode = 0
108
for file_group, changed_files in file_groups.items():
110
- follow_imports = None if ci and file_group == "" else "skip"
109
+ follow_imports = None if file_group == "" else "skip"
110
if changed_files:
111
returncode |= mypy(
112
changed_files, python_version, follow_imports, file_group
vllm/config/speculative.py
+38
-2
@@ -62,6 +62,7 @@ SpeculativeMethod = Literal[
62
"mlp_speculator",
63
"draft_model",
64
"suffix",
65
+ "custom_class",
66
EagleModelTypes,
67
NgramGPUTypes,
68
]
@@ -511,7 +512,16 @@ class SpeculativeConfig:
512
# default.
513
514
# infer method from user args
514
- if self.method is None:
515
+ # Check if the model field contains a custom module path (e.g., 'pkg.Mod')
516
+ if (
517
+ self.model is not None
518
+ and "." in self.model
519
+ and not self.model.startswith(("http://", "https://", "file://"))
520
+ and "/" not in self.model # not a HuggingFace repo (org/model)
521
+ ):
522
+ # Treat as a custom class path
523
+ self.method = "custom_class"
524
+ elif self.method is None:
525
if self.model in ("ngram", "[ngram]"):
526
self.method = "ngram"
527
else:
@@ -545,6 +555,14 @@ class SpeculativeConfig:
555
self.model = "suffix"
556
elif self.method == "extract_hidden_states":
557
self.model = "extract_hidden_states"
558
+ elif self.method == "custom_class":
559
+ # method was set explicitly, but model should already contain the
560
+ # custom module path. If not, this is a configuration error.
561
+ if self.model is None:
562
+ raise ValueError(
563
+ "method='custom_class' requires 'model' to contain the "
564
+ "custom proposer module path (e.g., 'my_module.MyProposer')."
565
+ )
566
else:
567
raise ValueError(
568
"num_speculative_tokens was provided but without speculative model."
@@ -588,6 +606,18 @@ class SpeculativeConfig:
606
self.draft_parallel_config = self.target_parallel_config
607
elif self.method == "suffix":
608
self._validate_suffix_decoding()
609
+ elif self.method == "custom_class":
610
+ # Custom class proposer does not need a draft model.
611
+ # It will dynamically load the user-provided class at runtime.
612
+ logger.warning_once(
613
+ "Using a custom class-based proposer backend. This is an "
614
+ "experimental feature and the proposer interface is subject to "
615
+ "breaking changes in future vLLM releases."
616
+ )
617
+ self.prompt_lookup_max = 0
618
+ self.prompt_lookup_min = 0
619
+ self.draft_model_config = self.target_model_config
620
+ self.draft_parallel_config = self.target_parallel_config
621
elif self.method == "extract_hidden_states":
622
from vllm.transformers_utils.configs.extract_hidden_states import (
623
ExtractHiddenStatesConfig,
@@ -1063,7 +1093,13 @@ class SpeculativeConfig:
1093
method = self.method
1094
model = (
1095
None
1066
- if method in ("ngram", "suffix", "extract_hidden_states")
1096
+ if method
1097
+ in (
1098
+ "ngram",
1099
+ "suffix",
1100
+ "extract_hidden_states",
1101
+ "custom_class",
1102
+ )
1103
else self.draft_model_config.model
1104
)
1105
num_spec_tokens = self.num_speculative_tokens
vllm/engine/arg_utils.py
-5
@@ -1613,11 +1613,6 @@ class EngineArgs:
1613
) -> SpeculativeConfig | None:
1614
"""Initializes and returns a SpeculativeConfig object based on
1615
`speculative_config`.
1616
-
1617
- This function utilizes `speculative_config` to create a
1618
- SpeculativeConfig object. The `speculative_config` can either be
1619
- provided as a JSON string input via CLI arguments or directly as a
1620
- dictionary from the engine.
1616
"""
1617
if self.speculative_config is None:
1618
return None
vllm/v1/spec_decode/custom_class_proposer.py
new
+73
@@ -0,0 +1,73 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.logger import init_logger
8
+
9
+logger = init_logger(__name__)
10
+
11
+
12
+def create_custom_proposer(vllm_config: VllmConfig):
13
+ """Load and instantiate a user-provided proposer class.
14
+
15
+ The class path is read from ``speculative_config.model``
16
+ (e.g., ``"my_module.MyCustomProposer"``). The class is
17
+ imported, instantiated with *vllm_config*, and returned
18
+ directly so the caller can use it without any wrapper.
19
+
20
+ The returned object must expose a callable ``propose`` method.
21
+ """
22
+ assert vllm_config.speculative_config is not None
23
+ spec_config = vllm_config.speculative_config
24
+
25
+ backend = spec_config.model
26
+ assert backend is not None
27
+
28
+ if "." not in backend:
29
+ raise ValueError(
30
+ f"Invalid custom proposer module path '{backend}'. "
31
+ "It must be a full module path (e.g., 'module.MyProposerClass')."
32
+ )
33
+
34
+ module_path, class_name = backend.rsplit(".", 1)
35
+ try:
36
+ module = importlib.import_module(module_path)
37
+ except ImportError as e:
38
+ raise ImportError(
39
+ f"Cannot import module '{module_path}' for custom proposer '{backend}': {e}"
40
+ ) from e
41
+
42
+ user_class = getattr(module, class_name, None)
43
+ if user_class is None:
44
+ raise AttributeError(
45
+ f"Module '{module_path}' has no attribute '{class_name}' "
46
+ f"(speculative_config.model='{backend}')"
47
+ )
48
+
49
+ try:
50
+ instance = user_class(vllm_config)
51
+ except Exception as e:
52
+ raise RuntimeError(
53
+ f"Failed to instantiate custom proposer class '{backend}': {e}. "
54
+ "The class constructor must accept VllmConfig as argument."
55
+ ) from e
56
+
57
+ if not hasattr(instance, "propose"):
58
+ raise AttributeError(
59
+ f"Custom proposer class '{backend}' must have a 'propose' method."
60
+ )
61
+ if not callable(instance.propose):
62
+ raise AttributeError(
63
+ f"Custom proposer class '{backend}' has a 'propose' attribute "
64
+ "but it is not callable."
65
+ )
66
+
67
+ logger.info(
68
+ "Loaded custom proposer class '%s' with num_speculative_tokens=%d",
69
+ backend,
70
+ spec_config.num_speculative_tokens,
71
+ )
72
+
73
+ return instance
vllm/v1/worker/gpu_model_runner.py
+16
-2
@@ -169,6 +169,7 @@ from vllm.v1.sample.logits_processor.interface import LogitsProcessor
169
from vllm.v1.sample.metadata import SamplingMetadata
170
from vllm.v1.sample.rejection_sampler import RejectionSampler
171
from vllm.v1.sample.sampler import Sampler
172
+from vllm.v1.spec_decode.custom_class_proposer import create_custom_proposer
173
from vllm.v1.spec_decode.dflash import DFlashProposer
174
from vllm.v1.spec_decode.draft_model import DraftModelProposer
175
from vllm.v1.spec_decode.eagle import EagleProposer
@@ -531,7 +532,11 @@ class GPUModelRunner(
532
| ExtractHiddenStatesProposer
533
| Gemma4Proposer
534
)
534
- if self.speculative_config.method == "ngram":
535
+ if self.speculative_config.method == "custom_class":
536
+ self.drafter = create_custom_proposer( # type: ignore[assignment]
537
+ self.vllm_config
538
+ )
539
+ elif self.speculative_config.method == "ngram":
540
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
541
542
self.drafter = NgramProposer(self.vllm_config)
@@ -4614,6 +4619,14 @@ class GPUModelRunner(
4619
self.input_batch.token_ids_cpu,
4620
slot_mappings=slot_mappings,
4621
)
4622
+ elif spec_config.method == "custom_class":
4623
+ assert isinstance(sampled_token_ids, list)
4624
+ draft_token_ids = cast(Any, self.drafter).propose(
4625
+ sampled_token_ids,
4626
+ self.input_batch.num_tokens_no_spec,
4627
+ self.input_batch.token_ids_cpu,
4628
+ slot_mappings=slot_mappings,
4629
+ )
4630
elif spec_config.use_ngram_gpu():
4631
assert isinstance(self.drafter, NgramProposerGPU)
4632
(
@@ -4885,7 +4898,8 @@ class GPUModelRunner(
4898
)
4899
if hasattr(self, "drafter"):
4900
logger.info_once("Loading drafter model...")
4888
- self.drafter.load_model(self.model)
4901
+ if hasattr(self.drafter, "load_model"):
4902
+ self.drafter.load_model(self.model)
4903
if (
4904
hasattr(self.drafter, "model")
4905
and is_mixture_of_experts(self.drafter.model)