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 14, 22:35 UTC
[Speculative Decoding] Add `speculators` config support (#21345)
Dipika Sikka committed
Aug 1, 2025 at 08:25 UTC
dfbc1f88807a1bddb75fc1dd587922567d7c133f
9 files changed
+232
-11
tests/speculative_decoding/speculators/test_eagle3.py
new
+16
@@ -0,0 +1,16 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+
7
+@pytest.mark.parametrize(
8
+ "model_path",
9
+ [("nm-testing/SpeculatorLlama3-1-8B-Eagle3-converted-0717"),
10
+ ("nm-testing/SpeculatorLlama3-1-8B-Eagle3-converted-0717-quantized")])
11
+def test_llama(vllm_runner, example_prompts, model_path):
12
+ with vllm_runner(model_path, dtype=torch.bfloat16) as vllm_model:
13
+ vllm_outputs = vllm_model.generate_greedy(example_prompts,
14
+ max_tokens=20)
15
+ print(vllm_outputs)
16
+ assert vllm_outputs
vllm/config.py
+16
-4
@@ -39,8 +39,8 @@ from vllm.transformers_utils.config import (
39
ConfigFormat, get_config, get_hf_image_processor_config,
40
get_hf_text_config, get_pooling_config,
41
get_sentence_transformer_tokenizer_config, is_encoder_decoder,
42
- try_get_generation_config, try_get_safetensors_metadata,
43
- try_get_tokenizer_config, uses_mrope)
42
+ maybe_override_with_speculators_target_model, try_get_generation_config,
43
+ try_get_safetensors_metadata, try_get_tokenizer_config, uses_mrope)
44
from vllm.transformers_utils.s3_utils import S3Model
45
from vllm.transformers_utils.utils import is_s3, maybe_model_redirect
46
# yapf conflicts with isort for this block
@@ -535,6 +535,15 @@ class ModelConfig:
535
"affect the random state of the Python process that "
536
"launched vLLM.", self.seed)
537
538
+ if self.runner != "draft":
539
+ # If we're not running the draft model, check for speculators config
540
+ # If speculators config, set model / tokenizer to be target model
541
+ self.model, self.tokenizer = maybe_override_with_speculators_target_model( # noqa: E501
542
+ model=self.model,
543
+ tokenizer=self.tokenizer,
544
+ revision=self.revision,
545
+ trust_remote_code=self.trust_remote_code)
546
+
547
# Keep set served_model_name before maybe_model_redirect(self.model)
548
self.served_model_name = get_served_model_name(self.model,
549
self.served_model_name)
@@ -606,8 +615,8 @@ class ModelConfig:
615
self.config_format,
616
hf_overrides_kw=hf_overrides_kw,
617
hf_overrides_fn=hf_overrides_fn)
609
- self.hf_config = hf_config
618
619
+ self.hf_config = hf_config
620
self.hf_text_config = get_hf_text_config(self.hf_config)
621
self.attention_chunk_size = getattr(self.hf_text_config,
622
"attention_chunk_size", None)
@@ -2980,10 +2989,13 @@ class SpeculativeConfig:
2989
"Chunked prefill and EAGLE are not compatible "
2990
"when using V0.")
2991
2992
+ from vllm.transformers_utils.configs import (
2993
+ SpeculatorsConfig)
2994
from vllm.transformers_utils.configs.eagle import (
2995
EAGLEConfig)
2996
+
2997
if isinstance(self.draft_model_config.hf_config,
2986
- EAGLEConfig):
2998
+ (EAGLEConfig, SpeculatorsConfig)):
2999
pass
3000
else:
3001
eagle_config = EAGLEConfig(
vllm/engine/arg_utils.py
+21
-1
@@ -978,8 +978,28 @@ class EngineArgs:
978
provided as a JSON string input via CLI arguments or directly as a
979
dictionary from the engine.
980
"""
981
+
982
+ from vllm.transformers_utils.config import get_config
983
+ from vllm.transformers_utils.configs.speculators.base import (
984
+ SpeculatorsConfig)
985
+
986
if self.speculative_config is None:
982
- return None
987
+ hf_config = get_config(self.hf_config_path or self.model,
988
+ self.trust_remote_code, self.revision,
989
+ self.code_revision, self.config_format)
990
+
991
+ # if loading a SpeculatorsConfig, load the specualtive_config
992
+ # details from the config directly
993
+ # no user input required / expected
994
+ if isinstance(hf_config, SpeculatorsConfig):
995
+ # We create one since we dont create one
996
+ self.speculative_config = {}
997
+ self.speculative_config[
998
+ "num_speculative_tokens"] = hf_config.num_lookahead_tokens
999
+ self.speculative_config["model"] = self.model
1000
+ self.speculative_config["method"] = hf_config.method
1001
+ else:
1002
+ return None
1003
1004
# Note(Shangming): These parameters are not obtained from the cli arg
1005
# '--speculative-config' and must be passed in when creating the engine
vllm/model_executor/models/llama_eagle3.py
+23
-3
@@ -51,6 +51,25 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
51
52
self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
53
54
+ if getattr(config, "norm_before_residual", False):
55
+ self._residual_norm = self._norm_before_residual
56
+ else:
57
+ self._residual_norm = self._norm_after_residual
58
+
59
+ def _norm_before_residual(
60
+ self,
61
+ hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
62
+ hidden_states = self.hidden_norm(hidden_states)
63
+ residual = hidden_states
64
+ return hidden_states, residual
65
+
66
+ def _norm_after_residual(
67
+ self,
68
+ hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
69
+ residual = hidden_states
70
+ hidden_states = self.hidden_norm(hidden_states)
71
+ return hidden_states, residual
72
+
73
def forward(
74
self,
75
positions: torch.Tensor,
@@ -59,9 +78,10 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
78
residual: Optional[torch.Tensor],
79
) -> tuple[torch.Tensor, torch.Tensor]:
80
62
- residual = hidden_states
81
embeds = self.input_layernorm(embeds)
64
- hidden_states = self.hidden_norm(hidden_states)
82
+
83
+ hidden_states, residual = self._residual_norm(
84
+ hidden_states=hidden_states)
85
86
hidden_states = torch.cat([embeds, hidden_states], dim=-1)
87
# Self Attention
@@ -102,7 +122,7 @@ class LlamaModel(nn.Module):
122
123
self.layers = nn.ModuleList([
124
LlamaDecoderLayer(
105
- self.config,
125
+ config=self.config,
126
prefix=maybe_prefix(prefix, f"layers.{start_layer_id}"),
127
)
128
])
vllm/transformers_utils/config.py
+29
-3
@@ -35,8 +35,9 @@ from vllm.transformers_utils.configs import (ChatGLMConfig, DeepseekVLV2Config,
35
MllamaConfig, MLPSpeculatorConfig,
36
Nemotron_Nano_VL_Config,
37
NemotronConfig, NVLM_D_Config,
38
- RWConfig, Step3TextConfig,
39
- Step3VLConfig, UltravoxConfig)
38
+ RWConfig, SpeculatorsConfig,
39
+ Step3TextConfig, Step3VLConfig,
40
+ UltravoxConfig)
41
# yapf: enable
42
from vllm.transformers_utils.configs.mistral import adapt_config_dict
43
from vllm.transformers_utils.utils import check_gguf_file
@@ -81,6 +82,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {
82
"mlp_speculator": MLPSpeculatorConfig,
83
"medusa": MedusaConfig,
84
"eagle": EAGLEConfig,
85
+ "speculators": SpeculatorsConfig,
86
"nemotron": NemotronConfig,
87
"NVLM_D": NVLM_D_Config,
88
"ultravox": UltravoxConfig,
@@ -287,6 +289,27 @@ def _maybe_remap_hf_config_attrs(config: PretrainedConfig) -> PretrainedConfig:
289
return config
290
291
292
+def maybe_override_with_speculators_target_model(
293
+ model: str,
294
+ tokenizer: str,
295
+ trust_remote_code: bool,
296
+ revision: Optional[str] = None) -> tuple[str, str]:
297
+ """
298
+ If running a speculators config, override running model with target model
299
+ """
300
+ config_dict, _ = PretrainedConfig.get_config_dict(
301
+ model,
302
+ revision=revision,
303
+ trust_remote_code=trust_remote_code,
304
+ token=_get_hf_token(),
305
+ )
306
+ spec_config = config_dict.get("speculators_config")
307
+ # Return the target model
308
+ if spec_config is not None:
309
+ model = tokenizer = spec_config["verifier"]["name_or_path"]
310
+ return model, tokenizer
311
+
312
+
313
def get_config(
314
model: Union[str, Path],
315
trust_remote_code: bool,
@@ -345,9 +368,12 @@ def get_config(
368
token=_get_hf_token(),
369
**kwargs,
370
)
348
-
371
# Use custom model class if it's in our registry
372
model_type = config_dict.get("model_type")
373
+ if model_type is None:
374
+ model_type = "speculators" if config_dict.get(
375
+ "speculators_config") is not None else model_type
376
+
377
if model_type in _CONFIG_REGISTRY:
378
config_class = _CONFIG_REGISTRY[model_type]
379
config = config_class.from_pretrained(
vllm/transformers_utils/configs/__init__.py
+2
@@ -24,6 +24,7 @@ from vllm.transformers_utils.configs.nemotron import NemotronConfig
24
from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
25
from vllm.transformers_utils.configs.nemotron_vl import Nemotron_Nano_VL_Config
26
from vllm.transformers_utils.configs.nvlm_d import NVLM_D_Config
27
+from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
28
from vllm.transformers_utils.configs.step3_vl import (Step3TextConfig,
29
Step3VisionEncoderConfig,
30
Step3VLConfig)
@@ -44,6 +45,7 @@ __all__ = [
45
"NemotronHConfig",
46
"Nemotron_Nano_VL_Config",
47
"NVLM_D_Config",
48
+ "SpeculatorsConfig",
49
"UltravoxConfig",
50
"Step3VLConfig",
51
"Step3VisionEncoderConfig",
vllm/transformers_utils/configs/speculators/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/transformers_utils/configs/speculators/algos.py
new
+32
@@ -0,0 +1,32 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+SUPPORTED_SPECULATORS_TYPES = {}
5
+
6
+
7
+def register_speculator(name):
8
+
9
+ def decorator(fn):
10
+ SUPPORTED_SPECULATORS_TYPES[name] = fn
11
+ return fn
12
+
13
+ return decorator
14
+
15
+
16
+@register_speculator("eagle3")
17
+def update_eagle3(config_dict: dict, vllm_config: dict) -> None:
18
+ """
19
+ Apply Eagle-3 specific configuration transformations.
20
+
21
+ Eagle-3 specific fields:
22
+ - draft_vocab_size: Size of the draft model's vocabulary
23
+ - target_hidden_size: Hidden size of the target model
24
+ - norm_before_residual: Whether to apply norm before residual connection
25
+ """
26
+
27
+ vllm_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
28
+ if config_dict.get("target_hidden_size") is not None:
29
+ vllm_config["target_hidden_size"] = config_dict["target_hidden_size"]
30
+ vllm_config["norm_before_residual"] = config_dict.get(
31
+ "norm_before_residual", True)
32
+ vllm_config["architectures"] = ["Eagle3LlamaForCausalLM"]
vllm/transformers_utils/configs/speculators/base.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
+from typing import Any, Union
5
+
6
+from transformers import PretrainedConfig
7
+
8
+from vllm.transformers_utils.configs.speculators.algos import (
9
+ SUPPORTED_SPECULATORS_TYPES)
10
+
11
+__all__ = ["SpeculatorsConfig"]
12
+
13
+
14
+class SpeculatorsConfig(PretrainedConfig):
15
+ model_type = "speculators"
16
+
17
+ @classmethod
18
+ def from_pretrained(
19
+ cls,
20
+ pretrained_model_name_or_path: Union[str, os.PathLike],
21
+ **kwargs,
22
+ ) -> "SpeculatorsConfig":
23
+ """Load speculators Eagle config and convert to vLLM format."""
24
+ config_dict, _ = cls.get_config_dict(pretrained_model_name_or_path,
25
+ **kwargs)
26
+
27
+ speculators_model_type = config_dict.get("speculators_model_type")
28
+ if speculators_model_type not in SUPPORTED_SPECULATORS_TYPES:
29
+ raise ValueError(
30
+ f"Expected one of: {SUPPORTED_SPECULATORS_TYPES}. "
31
+ "Please ensure you're loading a speculators-format model.")
32
+
33
+ # validate fields
34
+ # TODO: @dsikka - use speculators pydantic model to validate
35
+ cls.validate_speculators_config(config_dict=config_dict)
36
+ # Convert from speculators config -> format that can be ingested by vLLM
37
+ vllm_config = cls.convert_speculators_to_vllm(config_dict=config_dict)
38
+ # Apply anything specific to the supported algorithm
39
+ algo_updater = SUPPORTED_SPECULATORS_TYPES[speculators_model_type]
40
+ algo_updater(config_dict=config_dict, vllm_config=vllm_config)
41
+ return cls(**vllm_config)
42
+
43
+ @classmethod
44
+ def validate_speculators_config(cls, config_dict: dict[str, Any]) -> None:
45
+ try:
46
+ spec_config = config_dict["speculators_config"]
47
+ methods = spec_config["proposal_methods"]
48
+ first_method = methods[0]
49
+ _ = first_method["speculative_tokens"]
50
+ _ = spec_config["verifier"]["name_or_path"]
51
+ _ = config_dict["speculators_model_type"]
52
+ except (KeyError, IndexError, TypeError) as e:
53
+ raise ValueError("Invalid speculators config structure") from e
54
+
55
+ if "transformer_layer_config" not in config_dict:
56
+ raise ValueError("Must provide transformer_layer_config")
57
+
58
+ if not isinstance(config_dict["transformer_layer_config"], dict):
59
+ raise TypeError(
60
+ "'transformer_layer_config' must be a dictionary if provided")
61
+
62
+ @classmethod
63
+ def convert_speculators_to_vllm(
64
+ cls, config_dict: dict[str, Any]) -> dict[str, Any]:
65
+ """
66
+ Convert speculators config format to vLLM format.
67
+
68
+ This method handles the translation of field names and structure
69
+ between speculators and vLLM formats.
70
+
71
+ Returns:
72
+ Dictionary with vLLM-compatible configuration
73
+ """
74
+ # Currently we only support one proposal method
75
+ spec_config = config_dict["speculators_config"]
76
+ first_method = spec_config.get("proposal_methods")[0]
77
+ num_lookahead_tokens = first_method.get("speculative_tokens")
78
+
79
+ if num_lookahead_tokens is None:
80
+ raise ValueError(
81
+ "Missing 'speculative_tokens' in proposal method. "
82
+ f"Got: {first_method}")
83
+
84
+ # Build base vLLM config
85
+ vllm_config = {
86
+ "method": config_dict.get("speculators_model_type"),
87
+ "num_lookahead_tokens": num_lookahead_tokens,
88
+ "target_model": spec_config.get("verifier")["name_or_path"]
89
+ }
90
+ vllm_config.update(config_dict["transformer_layer_config"])
91
+ return vllm_config