@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0faf3cc3e8
Fork 0

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, 05:30 UTC

Move `SpeculativeConfig` from `config/__init__.py` to `config/speculative.py` (#24904)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Sep 16, 2025 at 12:51 UTC 0faf3cc3e84a83bac56492d297f2c9909fec7de9
2 files changed +556 -523
vllm/config/__init__.py
+2 -523
@@ -24,7 +24,7 @@ from pydantic import (ConfigDict, SkipValidation, field_validator,
24 model_validator)
25 from pydantic.dataclasses import dataclass
26 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
27 -from typing_extensions import Self, assert_never, runtime_checkable
27 +from typing_extensions import assert_never, runtime_checkable
28
29 import vllm.envs as envs
30 from vllm import version
@@ -41,6 +41,7 @@ from vllm.config.multimodal import (MMCacheType, MMEncoderTPMode,
41 from vllm.config.parallel import (DistributedExecutorBackend, EPLBConfig,
42 ParallelConfig)
43 from vllm.config.scheduler import SchedulerConfig, SchedulerPolicy
44 +from vllm.config.speculative import SpeculativeConfig
45 from vllm.config.utils import ConfigType, config
46 from vllm.logger import init_logger
47 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -1846,528 +1847,6 @@ class DeviceConfig:
1847 self.device = torch.device(self.device_type)
1848
1849
1849 -SpeculativeMethod = Literal["ngram", "eagle", "eagle3", "medusa",
1850 - "mlp_speculator", "draft_model", "deepseek_mtp",
1851 - "ernie_mtp", "qwen3_next_mtp"]
1852 -
1853 -
1854 -@config
1855 -@dataclass
1856 -class SpeculativeConfig:
1857 - """Configuration for speculative decoding."""
1858 -
1859 - # General speculative decoding control
1860 - num_speculative_tokens: SkipValidation[int] = None # type: ignore
1861 - """The number of speculative tokens, if provided. It will default to the
1862 - number in the draft model config if present, otherwise, it is required."""
1863 - model: Optional[str] = None
1864 - """The name of the draft model, eagle head, or additional weights, if
1865 - provided."""
1866 - method: Optional[SpeculativeMethod] = None
1867 - """The name of the speculative method to use. If users provide and set the
1868 - `model` param, the speculative method type will be detected automatically
1869 - if possible, if `model` param is not provided, the method name must be
1870 - provided.
1871 -
1872 - If using `ngram` method, the related configuration `prompt_lookup_max` and
1873 - `prompt_lookup_min` should be considered."""
1874 - draft_tensor_parallel_size: Optional[int] = None
1875 - """The degree of the tensor parallelism for the draft model. Can only be 1
1876 - or the same as the target model's tensor parallel size."""
1877 - disable_logprobs: bool = True
1878 - """If set to True, token log probabilities are not returned during
1879 - speculative decoding. If set to False, token log probabilities are returned
1880 - according to the log probability settings in SamplingParams."""
1881 -
1882 - # Draft model configuration
1883 - quantization: Optional[me_quant.QuantizationMethods] = None
1884 - """Quantization method that was used to quantize the draft model weights.
1885 - If `None`, we assume the model weights are not quantized. Note that it only
1886 - takes effect when using the draft model-based speculative method."""
1887 - max_model_len: Optional[int] = None
1888 - """The maximum model length of the draft model. Used when testing the
1889 - ability to skip speculation for some sequences."""
1890 - revision: Optional[str] = None
1891 - """The specific model version to use for the draft model. It can be a
1892 - branch name, a tag name, or a commit id. If unspecified, will use the
1893 - default version."""
1894 - code_revision: Optional[str] = None
1895 - """The specific revision to use for the draft model code on Hugging Face
1896 - Hub. It can be a branch name, a tag name, or a commit id. If unspecified,
1897 - will use the default version."""
1898 -
1899 - # Advanced control
1900 - disable_by_batch_size: Optional[int] = None
1901 - """Disable speculative decoding for new incoming requests when the number
1902 - of enqueued requests is larger than this value, if provided."""
1903 -
1904 - # Ngram proposer configuration
1905 - prompt_lookup_max: Optional[int] = None
1906 - """Maximum size of ngram token window when using Ngram proposer, required
1907 - when method is set to ngram."""
1908 - prompt_lookup_min: Optional[int] = None
1909 - """Minimum size of ngram token window when using Ngram proposer, if
1910 - provided. Defaults to 1."""
1911 -
1912 - speculative_token_tree: Optional[str] = None
1913 - """Specifies the tree structure for speculative token generation.
1914 - """
1915 - # required configuration params passed from engine
1916 - target_model_config: SkipValidation[ModelConfig] = None # type: ignore
1917 - """The configuration of the target model."""
1918 - target_parallel_config: SkipValidation[
1919 - ParallelConfig] = None # type: ignore
1920 - """The parallel configuration for the target model."""
1921 - enable_chunked_prefill: SkipValidation[bool] = None # type: ignore
1922 - """Whether vLLM is configured to use chunked prefill or not. Used for
1923 - raising an error since it's not yet compatible with speculative decode."""
1924 - disable_log_stats: SkipValidation[bool] = None # type: ignore
1925 - """Whether to disable the periodic printing of stage times in speculative
1926 - decoding."""
1927 -
1928 - # params generated in the post-init stage
1929 - draft_model_config: SkipValidation[ModelConfig] = None # type: ignore
1930 - """The configuration of the draft model initialized internal."""
1931 - draft_parallel_config: SkipValidation[
1932 - ParallelConfig] = None # type: ignore
1933 - """The parallel configuration for the draft model initialized internal."""
1934 -
1935 - def compute_hash(self) -> str:
1936 - """
1937 - WARNING: Whenever a new field is added to this config,
1938 - ensure that it is included in the factors list if
1939 - it affects the computation graph.
1940 -
1941 - Provide a hash that uniquely identifies all the configs
1942 - that affect the structure of the computation
1943 - graph from input ids/embeddings to the final hidden states,
1944 - excluding anything before input ids/embeddings and after
1945 - the final hidden states.
1946 - """
1947 - factors: list[Any] = []
1948 - # Eagle3 affects the computation graph because it returns intermediate
1949 - # hidden states in addition to the final hidden state.
1950 - factors.append(self.method == "eagle3")
1951 - hash_str = hashlib.md5(str(factors).encode(),
1952 - usedforsecurity=False).hexdigest()
1953 - return hash_str
1954 -
1955 - @staticmethod
1956 - def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
1957 - if hf_config.model_type == "deepseek_v3":
1958 - hf_config.model_type = "deepseek_mtp"
1959 - if hf_config.model_type == "deepseek_mtp":
1960 - n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
1961 - hf_config.update({
1962 - "n_predict": n_predict,
1963 - "architectures": ["DeepSeekMTPModel"]
1964 - })
1965 -
1966 - if hf_config.architectures[0] == "MiMoForCausalLM":
1967 - hf_config.model_type = "mimo_mtp"
1968 - n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
1969 - hf_config.update({
1970 - "num_hidden_layers": 0,
1971 - "n_predict": n_predict,
1972 - "architectures": ["MiMoMTPModel"]
1973 - })
1974 -
1975 - if hf_config.architectures[0] == "Glm4MoeForCausalLM":
1976 - hf_config.model_type = "glm4_moe_mtp"
1977 - n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
1978 - hf_config.update({
1979 - "num_hidden_layers": 0,
1980 - "n_predict": n_predict,
1981 - "architectures": ["Glm4MoeMTPModel"]
1982 - })
1983 -
1984 - if hf_config.model_type == "ernie4_5_moe":
1985 - hf_config.model_type = "ernie_mtp"
1986 - if hf_config.model_type == "ernie_mtp":
1987 - n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
1988 - hf_config.update({
1989 - "n_predict": n_predict,
1990 - "architectures": ["ErnieMTPModel"]
1991 - })
1992 -
1993 - if hf_config.model_type == "qwen3_next":
1994 - hf_config.model_type = "qwen3_next_mtp"
1995 - if hf_config.model_type == "qwen3_next_mtp":
1996 - n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
1997 - hf_config.update({
1998 - "n_predict": n_predict,
1999 - "architectures": ["Qwen3NextMTP"]
2000 - })
2001 -
2002 - return hf_config
2003 -
2004 - def __post_init__(self):
2005 -
2006 - # Note: "method" is a new parameter that helps to extend the
2007 - # configuration of non-model-based proposers, and the "model" parameter
2008 - # will be used to set the draft model, eagle head, or additional weight
2009 - # when needed. If users do not specify "method", the speculative method
2010 - # will be detected automatically if possible. If the speculative method
2011 - # can not be detected, it will be considered as the "draft_model" by
2012 - # default.
2013 -
2014 - if self.model is None and self.num_speculative_tokens is not None:
2015 - # TODO(Shangming): Refactor mtp configuration logic when supporting
2016 - # mtp acceleration for more models besides deepseek_v3
2017 - if self.target_model_config and \
2018 - (self.target_model_config.hf_text_config.model_type \
2019 - == "deepseek_v3" or
2020 - self.target_model_config.hf_text_config.model_type in
2021 - ("mimo","ernie4_5_moe", "qwen3_next")):
2022 - # use the draft model from the same model:
2023 - self.model = self.target_model_config.model
2024 - # Align the quantization of draft model for cases such as
2025 - # --quantization fp8 with a bf16 checkpoint.
2026 - if not self.quantization:
2027 - self.quantization = self.target_model_config.quantization
2028 - elif self.method in ("ngram", "[ngram]"):
2029 - self.model = "ngram"
2030 - else:
2031 - raise ValueError("num_speculative_tokens was provided without "
2032 - "speculative model.")
2033 -
2034 - # Automatically configure the method for ngram when "model" is used
2035 - # instead of "method"
2036 - if self.method is None and (self.model is not None
2037 - and self.model in ("ngram", "[ngram]")):
2038 - self.method = "ngram"
2039 -
2040 - if self.method in ("ngram", "[ngram]"):
2041 - # Unified to "ngram" internally
2042 - self.method = "ngram"
2043 - # Set default values if not provided
2044 - if (self.prompt_lookup_min is None
2045 - and self.prompt_lookup_max is None):
2046 - # TODO(woosuk): Tune these values. They are arbitrarily chosen.
2047 - self.prompt_lookup_min = 5
2048 - self.prompt_lookup_max = 5
2049 - elif self.prompt_lookup_min is None:
2050 - assert self.prompt_lookup_max is not None
2051 - self.prompt_lookup_min = self.prompt_lookup_max
2052 - elif self.prompt_lookup_max is None:
2053 - assert self.prompt_lookup_min is not None
2054 - self.prompt_lookup_max = self.prompt_lookup_min
2055 -
2056 - # Validate values
2057 - if self.prompt_lookup_min < 1:
2058 - raise ValueError(
2059 - f"prompt_lookup_min={self.prompt_lookup_min} must be > 0")
2060 - if self.prompt_lookup_max < 1:
2061 - raise ValueError(
2062 - f"prompt_lookup_max={self.prompt_lookup_max} must be > 0")
2063 - if self.prompt_lookup_min > self.prompt_lookup_max:
2064 - raise ValueError(
2065 - f"prompt_lookup_min={self.prompt_lookup_min} must "
2066 - f"be <= prompt_lookup_max={self.prompt_lookup_max}")
2067 -
2068 - # TODO: current we still need extract vocab_size from target model
2069 - # config, in future, we may try refactor it out, and set
2070 - # draft related config as None here.
2071 - self.draft_model_config = self.target_model_config
2072 - self.draft_parallel_config = self.target_parallel_config
2073 - else:
2074 - self.prompt_lookup_max = 0
2075 - self.prompt_lookup_min = 0
2076 -
2077 - if self.model is not None:
2078 - self.draft_model_config = ModelConfig(
2079 - model=self.model,
2080 - runner="draft",
2081 - tokenizer=self.target_model_config.tokenizer,
2082 - tokenizer_mode=self.target_model_config.tokenizer_mode,
2083 - trust_remote_code=self.target_model_config.
2084 - trust_remote_code,
2085 - allowed_local_media_path=self.target_model_config.
2086 - allowed_local_media_path,
2087 - dtype=self.target_model_config.dtype,
2088 - seed=self.target_model_config.seed,
2089 - revision=self.revision,
2090 - code_revision=self.code_revision,
2091 - tokenizer_revision=self.target_model_config.
2092 - tokenizer_revision,
2093 - spec_target_max_model_len=self.target_model_config.
2094 - max_model_len,
2095 - quantization=self.quantization,
2096 - enforce_eager=self.target_model_config.enforce_eager,
2097 - max_seq_len_to_capture=self.target_model_config.
2098 - max_seq_len_to_capture,
2099 - max_logprobs=self.target_model_config.max_logprobs,
2100 - hf_overrides=SpeculativeConfig.hf_config_override,
2101 - )
2102 -
2103 - # Automatically detect the method
2104 - if self.method in ('eagle', 'eagle3'):
2105 - pass
2106 - # examples:
2107 - # yuhuili/EAGLE-LLaMA3-Instruct-8B
2108 - # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B
2109 - # AngelSlim/Qwen3-8B_eagle3
2110 - elif "eagle-" in self.draft_model_config.model.lower():
2111 - self.method = "eagle"
2112 - elif "eagle3" in self.draft_model_config.model.lower():
2113 - self.method = "eagle3"
2114 - elif self.draft_model_config.hf_config.model_type == "medusa":
2115 - self.method = "medusa"
2116 - elif (self.draft_model_config.hf_config.model_type ==
2117 - "mlp_speculator"):
2118 - self.method = "mlp_speculator"
2119 - elif (self.draft_model_config.hf_config.model_type
2120 - in ("deepseek_mtp", "mimo_mtp", "glm4_moe_mtp")):
2121 - self.method = "deepseek_mtp"
2122 - if self.num_speculative_tokens > 1:
2123 - logger.warning(
2124 - "All Deepseek MTP models only have " \
2125 - "one layer. Might need some code changes " \
2126 - "to support multiple layers."
2127 - )
2128 - elif (self.draft_model_config.hf_config.model_type ==
2129 - "ernie_mtp"):
2130 - self.method = "ernie_mtp"
2131 - if self.num_speculative_tokens > 1:
2132 - logger.warning(
2133 - "All Ernie MTP models only have " \
2134 - "one layer. Might need some code changes " \
2135 - "to support multiple layers."
2136 - )
2137 - elif (self.draft_model_config.hf_config.model_type ==
2138 - "qwen3_next_mtp"):
2139 - self.method = "qwen3_next_mtp"
2140 - if self.num_speculative_tokens > 1:
2141 - logger.warning(
2142 - "All Qwen3Next MTP models only have " \
2143 - "one layer. Might need some code changes " \
2144 - "to support multiple layers."
2145 - )
2146 - else:
2147 - self.method = "draft_model"
2148 - raise NotImplementedError(
2149 - "Speculative decoding with draft model is not "
2150 - "supported yet. Please consider using other "
2151 - "speculative decoding methods such as ngram, medusa, "
2152 - "eagle, or deepseek_mtp.")
2153 -
2154 - # Replace hf_config for EAGLE draft_model
2155 - if self.method in ("eagle", "eagle3"):
2156 - if self.enable_chunked_prefill and not envs.VLLM_USE_V1:
2157 - raise ValueError(
2158 - "Chunked prefill and EAGLE are not compatible "
2159 - "when using V0.")
2160 -
2161 - from vllm.transformers_utils.configs import (
2162 - SpeculatorsConfig)
2163 - from vllm.transformers_utils.configs.eagle import (
2164 - EAGLEConfig)
2165 -
2166 - if isinstance(self.draft_model_config.hf_config,
2167 - (EAGLEConfig, SpeculatorsConfig)):
2168 - pass
2169 - else:
2170 - eagle_config = EAGLEConfig(
2171 - self.draft_model_config.hf_config,
2172 - method=self.method,
2173 - model_type="eagle")
2174 - self.draft_model_config.hf_config = eagle_config
2175 -
2176 - if (self.num_speculative_tokens is not None
2177 - and hasattr(self.draft_model_config.hf_config,
2178 - "num_lookahead_tokens")):
2179 - self.draft_model_config.hf_config.num_lookahead_tokens = \
2180 - self.num_speculative_tokens
2181 -
2182 - n_predict = getattr(self.draft_model_config.hf_config,
2183 - "n_predict", None)
2184 - if n_predict is not None:
2185 - if self.num_speculative_tokens is None:
2186 - # Default to max value defined in draft model config.
2187 - self.num_speculative_tokens = n_predict
2188 - elif self.num_speculative_tokens > n_predict and \
2189 - self.num_speculative_tokens % n_predict != 0:
2190 - # Ensure divisibility for MTP module reuse.
2191 - raise ValueError(
2192 - f"num_speculative_tokens:{self.num_speculative_tokens}"
2193 - f" must be divisible by {n_predict=}")
2194 -
2195 - if self.speculative_token_tree is None:
2196 - # Generate chain of tokens.
2197 - self.speculative_token_tree = str([
2198 - (i + 1) * (0, )
2199 - for i in range(self.num_speculative_tokens)
2200 - ])
2201 - else:
2202 - # Sort the token tree breadth-first.
2203 - tree_choices = ast.literal_eval(
2204 - self.speculative_token_tree)
2205 - self.speculative_token_tree = str(
2206 - sorted(tree_choices, key=lambda t: (len(t), t)))
2207 -
2208 - self.draft_tensor_parallel_size = \
2209 - SpeculativeConfig._verify_and_get_draft_tp(
2210 - self.target_parallel_config,
2211 - self.draft_tensor_parallel_size,
2212 - self.draft_model_config.hf_config
2213 - )
2214 -
2215 - self.draft_model_config.max_model_len = (
2216 - SpeculativeConfig._maybe_override_draft_max_model_len(
2217 - self.max_model_len,
2218 - self.draft_model_config.max_model_len,
2219 - self.target_model_config.max_model_len,
2220 - ))
2221 -
2222 - self.draft_parallel_config = (
2223 - SpeculativeConfig.create_draft_parallel_config(
2224 - self.target_parallel_config,
2225 - self.draft_tensor_parallel_size))
2226 -
2227 - @staticmethod
2228 - def _maybe_override_draft_max_model_len(
2229 - speculative_max_model_len: Optional[int],
2230 - draft_max_model_len: int,
2231 - target_max_model_len: int,
2232 - ) -> int:
2233 - """Determine the max sequence len for the draft model. This is usually
2234 - the draft_max_model_len, but may be the target_max_model_len if it is
2235 - less than the draft_max_model_len, or may be speculative_max_model_len
2236 - if it is specified.
2237 -
2238 - This is necessary so that sequences do not exceed the capacity of the
2239 - draft model or the target model.
2240 -
2241 - speculative_max_model_len is mainly used for testing that sequences can
2242 - skip speculation.
2243 - """
2244 -
2245 - if speculative_max_model_len is not None:
2246 -
2247 - if speculative_max_model_len > draft_max_model_len:
2248 - raise ValueError(f"{speculative_max_model_len=} cannot be "
2249 - f"larger than {draft_max_model_len=}")
2250 -
2251 - if speculative_max_model_len > target_max_model_len:
2252 - raise ValueError(f"{speculative_max_model_len=} cannot be "
2253 - f"larger than {target_max_model_len=}")
2254 -
2255 - return speculative_max_model_len
2256 -
2257 - return min(
2258 - draft_max_model_len,
2259 - target_max_model_len,
2260 - )
2261 -
2262 - @staticmethod
2263 - def _verify_and_get_draft_tp(
2264 - target_parallel_config: ParallelConfig,
2265 - speculative_draft_tensor_parallel_size: Optional[int],
2266 - draft_hf_config: PretrainedConfig) -> int:
2267 - """
2268 - Verifies and adjusts the tensor parallel size for a draft model
2269 - specified using speculative_draft_tensor_parallel_size.
2270 - """
2271 - # If speculative_draft_tensor_parallel_size is unset then set it
2272 - # appropriately else verify that it is set correctly.
2273 - if speculative_draft_tensor_parallel_size is None:
2274 - if draft_hf_config.model_type == "mlp_speculator":
2275 - speculative_draft_tensor_parallel_size = 1
2276 - if target_parallel_config.tensor_parallel_size > 1:
2277 - logger.warning(
2278 - "%s cannot currently be run with tp>1; "
2279 - "setting speculative_draft_tensor_parallel_size=1",
2280 - draft_hf_config.model_type)
2281 - else:
2282 - speculative_draft_tensor_parallel_size = \
2283 - target_parallel_config.tensor_parallel_size
2284 - elif speculative_draft_tensor_parallel_size not in (
2285 - 1, target_parallel_config.tensor_parallel_size):
2286 - raise ValueError(
2287 - f"{speculative_draft_tensor_parallel_size=} cannot be "
2288 - f"other value than 1 or target model tensor_parallel_size")
2289 - return speculative_draft_tensor_parallel_size
2290 -
2291 - @staticmethod
2292 - def create_draft_parallel_config(
2293 - target_parallel_config: ParallelConfig,
2294 - speculative_draft_tensor_parallel_size: int,
2295 - ) -> ParallelConfig:
2296 - """Create a parallel config for use by the draft worker.
2297 -
2298 - This is mostly a copy of the target parallel config, except the tp_size.
2299 - """
2300 - draft_parallel_config = ParallelConfig(
2301 - pipeline_parallel_size=target_parallel_config.
2302 - pipeline_parallel_size,
2303 - tensor_parallel_size=speculative_draft_tensor_parallel_size,
2304 - distributed_executor_backend=target_parallel_config.
2305 - distributed_executor_backend,
2306 - max_parallel_loading_workers=target_parallel_config.
2307 - max_parallel_loading_workers,
2308 - disable_custom_all_reduce=target_parallel_config.
2309 - disable_custom_all_reduce,
2310 - ray_workers_use_nsight=target_parallel_config.
2311 - ray_workers_use_nsight,
2312 - placement_group=target_parallel_config.placement_group,
2313 - )
2314 -
2315 - return draft_parallel_config
2316 -
2317 - @model_validator(mode='after')
2318 - def _verify_args(self) -> Self:
2319 - if self.num_speculative_tokens is None:
2320 - raise ValueError(
2321 - "num_speculative_tokens must be provided with "
2322 - "speculative model unless the draft model config contains an "
2323 - "n_predict parameter.")
2324 -
2325 - if self.num_speculative_tokens <= 0:
2326 - raise ValueError("Expected num_speculative_tokens to be greater "
2327 - f"than zero ({self.num_speculative_tokens}).")
2328 -
2329 - if self.draft_model_config:
2330 - self.draft_model_config.verify_with_parallel_config(
2331 - self.draft_parallel_config)
2332 -
2333 - if (self.disable_by_batch_size is not None
2334 - and self.disable_by_batch_size < 2):
2335 - raise ValueError("Expect the batch size threshold of disabling "
2336 - "speculative decoding is > 1, but got "
2337 - f"{self.disable_by_batch_size=}")
2338 -
2339 - eagle3_target_supported = ["llama", "qwen"]
2340 - if self.method == "eagle3" and self.target_model_config and not any(
2341 - supported_model in
2342 - self.target_model_config.hf_text_config.model_type
2343 - for supported_model in eagle3_target_supported):
2344 - raise ValueError(
2345 - f"Eagle3 is only supported for {eagle3_target_supported} models. " # noqa: E501
2346 - f"Got {self.target_model_config.hf_text_config.model_type=}")
2347 -
2348 - return self
2349 -
2350 - @property
2351 - def num_lookahead_slots(self) -> int:
2352 - """The number of additional slots the scheduler should allocate per
2353 - step, in addition to the slots allocated for each known token.
2354 -
2355 - This is equal to the number of speculative tokens, as each speculative
2356 - token must be scored.
2357 - """
2358 - return self.num_speculative_tokens
2359 -
2360 - def use_eagle(self) -> bool:
2361 - return self.method in ("eagle", "eagle3", "deepseek_mtp", "ernie_mtp",
2362 - "qwen3_next_mtp")
2363 -
2364 - def __repr__(self) -> str:
2365 - method = self.method
2366 - model = None if method == "ngram" else self.draft_model_config.model
2367 - num_spec_tokens = self.num_speculative_tokens
2368 - return f"SpeculativeConfig({method=}, {model=}, {num_spec_tokens=})"
2369 -
2370 -
1850 @config
1851 @dataclass
1852 class PoolerConfig:
vllm/config/speculative.py new
+554
@@ -0,0 +1,554 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import ast
5 +import hashlib
6 +from typing import TYPE_CHECKING, Any, Literal, Optional
7 +
8 +from pydantic import SkipValidation, model_validator
9 +from pydantic.dataclasses import dataclass
10 +from typing_extensions import Self
11 +
12 +import vllm.envs as envs
13 +from vllm.config.parallel import ParallelConfig
14 +from vllm.config.utils import config
15 +from vllm.logger import init_logger
16 +from vllm.utils import LazyLoader
17 +
18 +if TYPE_CHECKING:
19 + from transformers import PretrainedConfig
20 +
21 + import vllm.model_executor.layers.quantization as me_quant
22 + from vllm.config import ModelConfig
23 +else:
24 + PretrainedConfig = Any
25 + ModelConfig = Any
26 +
27 + me_quant = LazyLoader("model_executor", globals(),
28 + "vllm.model_executor.layers.quantization")
29 +
30 +logger = init_logger(__name__)
31 +
32 +SpeculativeMethod = Literal["ngram", "eagle", "eagle3", "medusa",
33 + "mlp_speculator", "draft_model", "deepseek_mtp",
34 + "ernie_mtp", "qwen3_next_mtp"]
35 +
36 +
37 +@config
38 +@dataclass
39 +class SpeculativeConfig:
40 + """Configuration for speculative decoding."""
41 +
42 + # General speculative decoding control
43 + num_speculative_tokens: SkipValidation[int] = None # type: ignore
44 + """The number of speculative tokens, if provided. It will default to the
45 + number in the draft model config if present, otherwise, it is required."""
46 + model: Optional[str] = None
47 + """The name of the draft model, eagle head, or additional weights, if
48 + provided."""
49 + method: Optional[SpeculativeMethod] = None
50 + """The name of the speculative method to use. If users provide and set the
51 + `model` param, the speculative method type will be detected automatically
52 + if possible, if `model` param is not provided, the method name must be
53 + provided.
54 +
55 + If using `ngram` method, the related configuration `prompt_lookup_max` and
56 + `prompt_lookup_min` should be considered."""
57 + draft_tensor_parallel_size: Optional[int] = None
58 + """The degree of the tensor parallelism for the draft model. Can only be 1
59 + or the same as the target model's tensor parallel size."""
60 + disable_logprobs: bool = True
61 + """If set to True, token log probabilities are not returned during
62 + speculative decoding. If set to False, token log probabilities are returned
63 + according to the log probability settings in SamplingParams."""
64 +
65 + # Draft model configuration
66 + quantization: Optional[me_quant.QuantizationMethods] = None
67 + """Quantization method that was used to quantize the draft model weights.
68 + If `None`, we assume the model weights are not quantized. Note that it only
69 + takes effect when using the draft model-based speculative method."""
70 + max_model_len: Optional[int] = None
71 + """The maximum model length of the draft model. Used when testing the
72 + ability to skip speculation for some sequences."""
73 + revision: Optional[str] = None
74 + """The specific model version to use for the draft model. It can be a
75 + branch name, a tag name, or a commit id. If unspecified, will use the
76 + default version."""
77 + code_revision: Optional[str] = None
78 + """The specific revision to use for the draft model code on Hugging Face
79 + Hub. It can be a branch name, a tag name, or a commit id. If unspecified,
80 + will use the default version."""
81 +
82 + # Advanced control
83 + disable_by_batch_size: Optional[int] = None
84 + """Disable speculative decoding for new incoming requests when the number
85 + of enqueued requests is larger than this value, if provided."""
86 +
87 + # Ngram proposer configuration
88 + prompt_lookup_max: Optional[int] = None
89 + """Maximum size of ngram token window when using Ngram proposer, required
90 + when method is set to ngram."""
91 + prompt_lookup_min: Optional[int] = None
92 + """Minimum size of ngram token window when using Ngram proposer, if
93 + provided. Defaults to 1."""
94 +
95 + speculative_token_tree: Optional[str] = None
96 + """Specifies the tree structure for speculative token generation.
97 + """
98 + # required configuration params passed from engine
99 + target_model_config: SkipValidation[ModelConfig] = None # type: ignore
100 + """The configuration of the target model."""
101 + target_parallel_config: SkipValidation[
102 + ParallelConfig] = None # type: ignore
103 + """The parallel configuration for the target model."""
104 + enable_chunked_prefill: SkipValidation[bool] = None # type: ignore
105 + """Whether vLLM is configured to use chunked prefill or not. Used for
106 + raising an error since it's not yet compatible with speculative decode."""
107 + disable_log_stats: SkipValidation[bool] = None # type: ignore
108 + """Whether to disable the periodic printing of stage times in speculative
109 + decoding."""
110 +
111 + # params generated in the post-init stage
112 + draft_model_config: SkipValidation[ModelConfig] = None # type: ignore
113 + """The configuration of the draft model initialized internal."""
114 + draft_parallel_config: SkipValidation[
115 + ParallelConfig] = None # type: ignore
116 + """The parallel configuration for the draft model initialized internal."""
117 +
118 + def compute_hash(self) -> str:
119 + """
120 + WARNING: Whenever a new field is added to this config,
121 + ensure that it is included in the factors list if
122 + it affects the computation graph.
123 +
124 + Provide a hash that uniquely identifies all the configs
125 + that affect the structure of the computation
126 + graph from input ids/embeddings to the final hidden states,
127 + excluding anything before input ids/embeddings and after
128 + the final hidden states.
129 + """
130 + factors: list[Any] = []
131 + # Eagle3 affects the computation graph because it returns intermediate
132 + # hidden states in addition to the final hidden state.
133 + factors.append(self.method == "eagle3")
134 + hash_str = hashlib.md5(str(factors).encode(),
135 + usedforsecurity=False).hexdigest()
136 + return hash_str
137 +
138 + @staticmethod
139 + def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
140 + if hf_config.model_type == "deepseek_v3":
141 + hf_config.model_type = "deepseek_mtp"
142 + if hf_config.model_type == "deepseek_mtp":
143 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
144 + hf_config.update({
145 + "n_predict": n_predict,
146 + "architectures": ["DeepSeekMTPModel"]
147 + })
148 +
149 + if hf_config.architectures[0] == "MiMoForCausalLM":
150 + hf_config.model_type = "mimo_mtp"
151 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
152 + hf_config.update({
153 + "num_hidden_layers": 0,
154 + "n_predict": n_predict,
155 + "architectures": ["MiMoMTPModel"]
156 + })
157 +
158 + if hf_config.architectures[0] == "Glm4MoeForCausalLM":
159 + hf_config.model_type = "glm4_moe_mtp"
160 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
161 + hf_config.update({
162 + "num_hidden_layers": 0,
163 + "n_predict": n_predict,
164 + "architectures": ["Glm4MoeMTPModel"]
165 + })
166 +
167 + if hf_config.model_type == "ernie4_5_moe":
168 + hf_config.model_type = "ernie_mtp"
169 + if hf_config.model_type == "ernie_mtp":
170 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
171 + hf_config.update({
172 + "n_predict": n_predict,
173 + "architectures": ["ErnieMTPModel"]
174 + })
175 +
176 + if hf_config.model_type == "qwen3_next":
177 + hf_config.model_type = "qwen3_next_mtp"
178 + if hf_config.model_type == "qwen3_next_mtp":
179 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
180 + hf_config.update({
181 + "n_predict": n_predict,
182 + "architectures": ["Qwen3NextMTP"]
183 + })
184 +
185 + return hf_config
186 +
187 + def __post_init__(self):
188 +
189 + # Note: "method" is a new parameter that helps to extend the
190 + # configuration of non-model-based proposers, and the "model" parameter
191 + # will be used to set the draft model, eagle head, or additional weight
192 + # when needed. If users do not specify "method", the speculative method
193 + # will be detected automatically if possible. If the speculative method
194 + # can not be detected, it will be considered as the "draft_model" by
195 + # default.
196 +
197 + if self.model is None and self.num_speculative_tokens is not None:
198 + # TODO(Shangming): Refactor mtp configuration logic when supporting
199 + # mtp acceleration for more models besides deepseek_v3
200 + if self.target_model_config and \
201 + (self.target_model_config.hf_text_config.model_type \
202 + == "deepseek_v3" or
203 + self.target_model_config.hf_text_config.model_type in
204 + ("mimo","ernie4_5_moe", "qwen3_next")):
205 + # use the draft model from the same model:
206 + self.model = self.target_model_config.model
207 + # Align the quantization of draft model for cases such as
208 + # --quantization fp8 with a bf16 checkpoint.
209 + if not self.quantization:
210 + self.quantization = self.target_model_config.quantization
211 + elif self.method in ("ngram", "[ngram]"):
212 + self.model = "ngram"
213 + else:
214 + raise ValueError("num_speculative_tokens was provided without "
215 + "speculative model.")
216 +
217 + # Automatically configure the method for ngram when "model" is used
218 + # instead of "method"
219 + if self.method is None and (self.model is not None
220 + and self.model in ("ngram", "[ngram]")):
221 + self.method = "ngram"
222 +
223 + if self.method in ("ngram", "[ngram]"):
224 + # Unified to "ngram" internally
225 + self.method = "ngram"
226 + # Set default values if not provided
227 + if (self.prompt_lookup_min is None
228 + and self.prompt_lookup_max is None):
229 + # TODO(woosuk): Tune these values. They are arbitrarily chosen.
230 + self.prompt_lookup_min = 5
231 + self.prompt_lookup_max = 5
232 + elif self.prompt_lookup_min is None:
233 + assert self.prompt_lookup_max is not None
234 + self.prompt_lookup_min = self.prompt_lookup_max
235 + elif self.prompt_lookup_max is None:
236 + assert self.prompt_lookup_min is not None
237 + self.prompt_lookup_max = self.prompt_lookup_min
238 +
239 + # Validate values
240 + if self.prompt_lookup_min < 1:
241 + raise ValueError(
242 + f"prompt_lookup_min={self.prompt_lookup_min} must be > 0")
243 + if self.prompt_lookup_max < 1:
244 + raise ValueError(
245 + f"prompt_lookup_max={self.prompt_lookup_max} must be > 0")
246 + if self.prompt_lookup_min > self.prompt_lookup_max:
247 + raise ValueError(
248 + f"prompt_lookup_min={self.prompt_lookup_min} must "
249 + f"be <= prompt_lookup_max={self.prompt_lookup_max}")
250 +
251 + # TODO: current we still need extract vocab_size from target model
252 + # config, in future, we may try refactor it out, and set
253 + # draft related config as None here.
254 + self.draft_model_config = self.target_model_config
255 + self.draft_parallel_config = self.target_parallel_config
256 + else:
257 + self.prompt_lookup_max = 0
258 + self.prompt_lookup_min = 0
259 +
260 + if self.model is not None:
261 + # TODO: Move this import to the top once `ModelConfig`
262 + # lives in `vllm.config.model`.
263 + from vllm.config import ModelConfig
264 + self.draft_model_config = ModelConfig(
265 + model=self.model,
266 + runner="draft",
267 + tokenizer=self.target_model_config.tokenizer,
268 + tokenizer_mode=self.target_model_config.tokenizer_mode,
269 + trust_remote_code=self.target_model_config.
270 + trust_remote_code,
271 + allowed_local_media_path=self.target_model_config.
272 + allowed_local_media_path,
273 + dtype=self.target_model_config.dtype,
274 + seed=self.target_model_config.seed,
275 + revision=self.revision,
276 + code_revision=self.code_revision,
277 + tokenizer_revision=self.target_model_config.
278 + tokenizer_revision,
279 + spec_target_max_model_len=self.target_model_config.
280 + max_model_len,
281 + quantization=self.quantization,
282 + enforce_eager=self.target_model_config.enforce_eager,
283 + max_seq_len_to_capture=self.target_model_config.
284 + max_seq_len_to_capture,
285 + max_logprobs=self.target_model_config.max_logprobs,
286 + hf_overrides=SpeculativeConfig.hf_config_override,
287 + )
288 +
289 + # Automatically detect the method
290 + if self.method in ('eagle', 'eagle3'):
291 + pass
292 + # examples:
293 + # yuhuili/EAGLE-LLaMA3-Instruct-8B
294 + # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B
295 + # AngelSlim/Qwen3-8B_eagle3
296 + elif "eagle-" in self.draft_model_config.model.lower():
297 + self.method = "eagle"
298 + elif "eagle3" in self.draft_model_config.model.lower():
299 + self.method = "eagle3"
300 + elif self.draft_model_config.hf_config.model_type == "medusa":
301 + self.method = "medusa"
302 + elif (self.draft_model_config.hf_config.model_type ==
303 + "mlp_speculator"):
304 + self.method = "mlp_speculator"
305 + elif (self.draft_model_config.hf_config.model_type
306 + in ("deepseek_mtp", "mimo_mtp", "glm4_moe_mtp")):
307 + self.method = "deepseek_mtp"
308 + if self.num_speculative_tokens > 1:
309 + logger.warning(
310 + "All Deepseek MTP models only have " \
311 + "one layer. Might need some code changes " \
312 + "to support multiple layers."
313 + )
314 + elif (self.draft_model_config.hf_config.model_type ==
315 + "ernie_mtp"):
316 + self.method = "ernie_mtp"
317 + if self.num_speculative_tokens > 1:
318 + logger.warning(
319 + "All Ernie MTP models only have " \
320 + "one layer. Might need some code changes " \
321 + "to support multiple layers."
322 + )
323 + elif (self.draft_model_config.hf_config.model_type ==
324 + "qwen3_next_mtp"):
325 + self.method = "qwen3_next_mtp"
326 + if self.num_speculative_tokens > 1:
327 + logger.warning(
328 + "All Qwen3Next MTP models only have " \
329 + "one layer. Might need some code changes " \
330 + "to support multiple layers."
331 + )
332 + else:
333 + self.method = "draft_model"
334 + raise NotImplementedError(
335 + "Speculative decoding with draft model is not "
336 + "supported yet. Please consider using other "
337 + "speculative decoding methods such as ngram, medusa, "
338 + "eagle, or deepseek_mtp.")
339 +
340 + # Replace hf_config for EAGLE draft_model
341 + if self.method in ("eagle", "eagle3"):
342 + if self.enable_chunked_prefill and not envs.VLLM_USE_V1:
343 + raise ValueError(
344 + "Chunked prefill and EAGLE are not compatible "
345 + "when using V0.")
346 +
347 + from vllm.transformers_utils.configs import (
348 + SpeculatorsConfig)
349 + from vllm.transformers_utils.configs.eagle import (
350 + EAGLEConfig)
351 +
352 + if isinstance(self.draft_model_config.hf_config,
353 + (EAGLEConfig, SpeculatorsConfig)):
354 + pass
355 + else:
356 + eagle_config = EAGLEConfig(
357 + self.draft_model_config.hf_config,
358 + method=self.method,
359 + model_type="eagle")
360 + self.draft_model_config.hf_config = eagle_config
361 +
362 + if (self.num_speculative_tokens is not None
363 + and hasattr(self.draft_model_config.hf_config,
364 + "num_lookahead_tokens")):
365 + self.draft_model_config.hf_config.num_lookahead_tokens = \
366 + self.num_speculative_tokens
367 +
368 + n_predict = getattr(self.draft_model_config.hf_config,
369 + "n_predict", None)
370 + if n_predict is not None:
371 + if self.num_speculative_tokens is None:
372 + # Default to max value defined in draft model config.
373 + self.num_speculative_tokens = n_predict
374 + elif self.num_speculative_tokens > n_predict and \
375 + self.num_speculative_tokens % n_predict != 0:
376 + # Ensure divisibility for MTP module reuse.
377 + raise ValueError(
378 + f"num_speculative_tokens:{self.num_speculative_tokens}"
379 + f" must be divisible by {n_predict=}")
380 +
381 + if self.speculative_token_tree is None:
382 + # Generate chain of tokens.
383 + self.speculative_token_tree = str([
384 + (i + 1) * (0, )
385 + for i in range(self.num_speculative_tokens)
386 + ])
387 + else:
388 + # Sort the token tree breadth-first.
389 + tree_choices = ast.literal_eval(
390 + self.speculative_token_tree)
391 + self.speculative_token_tree = str(
392 + sorted(tree_choices, key=lambda t: (len(t), t)))
393 +
394 + self.draft_tensor_parallel_size = \
395 + SpeculativeConfig._verify_and_get_draft_tp(
396 + self.target_parallel_config,
397 + self.draft_tensor_parallel_size,
398 + self.draft_model_config.hf_config
399 + )
400 +
401 + self.draft_model_config.max_model_len = (
402 + SpeculativeConfig._maybe_override_draft_max_model_len(
403 + self.max_model_len,
404 + self.draft_model_config.max_model_len,
405 + self.target_model_config.max_model_len,
406 + ))
407 +
408 + self.draft_parallel_config = (
409 + SpeculativeConfig.create_draft_parallel_config(
410 + self.target_parallel_config,
411 + self.draft_tensor_parallel_size))
412 +
413 + @staticmethod
414 + def _maybe_override_draft_max_model_len(
415 + speculative_max_model_len: Optional[int],
416 + draft_max_model_len: int,
417 + target_max_model_len: int,
418 + ) -> int:
419 + """Determine the max sequence len for the draft model. This is usually
420 + the draft_max_model_len, but may be the target_max_model_len if it is
421 + less than the draft_max_model_len, or may be speculative_max_model_len
422 + if it is specified.
423 +
424 + This is necessary so that sequences do not exceed the capacity of the
425 + draft model or the target model.
426 +
427 + speculative_max_model_len is mainly used for testing that sequences can
428 + skip speculation.
429 + """
430 +
431 + if speculative_max_model_len is not None:
432 +
433 + if speculative_max_model_len > draft_max_model_len:
434 + raise ValueError(f"{speculative_max_model_len=} cannot be "
435 + f"larger than {draft_max_model_len=}")
436 +
437 + if speculative_max_model_len > target_max_model_len:
438 + raise ValueError(f"{speculative_max_model_len=} cannot be "
439 + f"larger than {target_max_model_len=}")
440 +
441 + return speculative_max_model_len
442 +
443 + return min(
444 + draft_max_model_len,
445 + target_max_model_len,
446 + )
447 +
448 + @staticmethod
449 + def _verify_and_get_draft_tp(
450 + target_parallel_config: ParallelConfig,
451 + speculative_draft_tensor_parallel_size: Optional[int],
452 + draft_hf_config: PretrainedConfig) -> int:
453 + """
454 + Verifies and adjusts the tensor parallel size for a draft model
455 + specified using speculative_draft_tensor_parallel_size.
456 + """
457 + # If speculative_draft_tensor_parallel_size is unset then set it
458 + # appropriately else verify that it is set correctly.
459 + if speculative_draft_tensor_parallel_size is None:
460 + if draft_hf_config.model_type == "mlp_speculator":
461 + speculative_draft_tensor_parallel_size = 1
462 + if target_parallel_config.tensor_parallel_size > 1:
463 + logger.warning(
464 + "%s cannot currently be run with tp>1; "
465 + "setting speculative_draft_tensor_parallel_size=1",
466 + draft_hf_config.model_type)
467 + else:
468 + speculative_draft_tensor_parallel_size = \
469 + target_parallel_config.tensor_parallel_size
470 + elif speculative_draft_tensor_parallel_size not in (
471 + 1, target_parallel_config.tensor_parallel_size):
472 + raise ValueError(
473 + f"{speculative_draft_tensor_parallel_size=} cannot be "
474 + f"other value than 1 or target model tensor_parallel_size")
475 + return speculative_draft_tensor_parallel_size
476 +
477 + @staticmethod
478 + def create_draft_parallel_config(
479 + target_parallel_config: ParallelConfig,
480 + speculative_draft_tensor_parallel_size: int,
481 + ) -> ParallelConfig:
482 + """Create a parallel config for use by the draft worker.
483 +
484 + This is mostly a copy of the target parallel config, except the tp_size.
485 + """
486 + draft_parallel_config = ParallelConfig(
487 + pipeline_parallel_size=target_parallel_config.
488 + pipeline_parallel_size,
489 + tensor_parallel_size=speculative_draft_tensor_parallel_size,
490 + distributed_executor_backend=target_parallel_config.
491 + distributed_executor_backend,
492 + max_parallel_loading_workers=target_parallel_config.
493 + max_parallel_loading_workers,
494 + disable_custom_all_reduce=target_parallel_config.
495 + disable_custom_all_reduce,
496 + ray_workers_use_nsight=target_parallel_config.
497 + ray_workers_use_nsight,
498 + placement_group=target_parallel_config.placement_group,
499 + )
500 +
501 + return draft_parallel_config
502 +
503 + @model_validator(mode='after')
504 + def _verify_args(self) -> Self:
505 + if self.num_speculative_tokens is None:
506 + raise ValueError(
507 + "num_speculative_tokens must be provided with "
508 + "speculative model unless the draft model config contains an "
509 + "n_predict parameter.")
510 +
511 + if self.num_speculative_tokens <= 0:
512 + raise ValueError("Expected num_speculative_tokens to be greater "
513 + f"than zero ({self.num_speculative_tokens}).")
514 +
515 + if self.draft_model_config:
516 + self.draft_model_config.verify_with_parallel_config(
517 + self.draft_parallel_config)
518 +
519 + if (self.disable_by_batch_size is not None
520 + and self.disable_by_batch_size < 2):
521 + raise ValueError("Expect the batch size threshold of disabling "
522 + "speculative decoding is > 1, but got "
523 + f"{self.disable_by_batch_size=}")
524 +
525 + eagle3_target_supported = ["llama", "qwen"]
526 + if self.method == "eagle3" and self.target_model_config and not any(
527 + supported_model in
528 + self.target_model_config.hf_text_config.model_type
529 + for supported_model in eagle3_target_supported):
530 + raise ValueError(
531 + f"Eagle3 is only supported for {eagle3_target_supported} models. " # noqa: E501
532 + f"Got {self.target_model_config.hf_text_config.model_type=}")
533 +
534 + return self
535 +
536 + @property
537 + def num_lookahead_slots(self) -> int:
538 + """The number of additional slots the scheduler should allocate per
539 + step, in addition to the slots allocated for each known token.
540 +
541 + This is equal to the number of speculative tokens, as each speculative
542 + token must be scored.
543 + """
544 + return self.num_speculative_tokens
545 +
546 + def use_eagle(self) -> bool:
547 + return self.method in ("eagle", "eagle3", "deepseek_mtp", "ernie_mtp",
548 + "qwen3_next_mtp")
549 +
550 + def __repr__(self) -> str:
551 + method = self.method
552 + model = None if method == "ngram" else self.draft_model_config.model
553 + num_spec_tokens = self.num_speculative_tokens
554 + return f"SpeculativeConfig({method=}, {model=}, {num_spec_tokens=})"