@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f525c0be8b
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 15, 22:15 UTC

[Model][Speculative Decoding] DeepSeek MTP spec decode (#12755)

Signed-off-by: Lu Fang <fanglu@fb.com> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>

Lucia Fang committed Feb 19, 2025 at 01:06 UTC f525c0be8b4c9ede079f866721526eba4c679ee3
14 files changed +727 -46
.buildkite/test-pipeline.yaml
+11 -11
@@ -2,7 +2,7 @@
2 # adding a new command to an existing step. See different options here for examples.
3
4 # This script will be feed into Jinja template in `test-template-aws.j2` at
5 -# https://github.com/vllm-project/buildkite-ci/blob/main/scripts/test-template-aws.j2
5 +# https://github.com/vllm-project/buildkite-ci/blob/main/scripts/test-template-aws.j2
6 # to generate the final pipeline yaml file.
7
8 # Documentation
@@ -15,7 +15,7 @@
15 # mirror_hardwares(list): the list of hardwares to run the test on as well. currently only supports [amd]
16 # gpu(str): override the GPU selection for the test. default is on L4 GPUs. currently only supports a100
17 # num_gpus(int): override the number of GPUs for the test. default to 1 GPU. currently support 2,4.
18 -# num_nodes(int): whether to simulate multi-node setup by launch multiple containers on one host,
18 +# num_nodes(int): whether to simulate multi-node setup by launch multiple containers on one host,
19 # in this case, commands must be specified. the first command runs on first host, the second
20 # command runs on the second host.
21 # working_dir(str): specify the place where command should execute, default to /vllm-workspace/tests
@@ -24,8 +24,8 @@
24 # When adding a test
25 # - If the test belong to an existing group, add it there
26 # - If the test is short, add to any existing step
27 -# - If the test takes more than 10min, then it is okay to create a new step.
28 -# Note that all steps execute in parallel.
27 +# - If the test takes more than 10min, then it is okay to create a new step.
28 +# Note that all steps execute in parallel.
29
30 steps:
31 ##### fast check tests #####
@@ -145,14 +145,14 @@ steps:
145 - RAY_DEDUP_LOGS=0 python3 ../examples/offline_inference/rlhf_colocate.py
146
147 - label: Metrics, Tracing Test # 10min
148 - num_gpus: 2
148 + num_gpus: 2
149 fast_check: true
150 source_file_dependencies:
151 - vllm/
152 - tests/metrics
153 - tests/tracing
154 commands:
155 - - pytest -v -s metrics
155 + - pytest -v -s metrics
156 - "pip install \
157 'opentelemetry-sdk>=1.26.0,<1.27.0' \
158 'opentelemetry-api>=1.26.0,<1.27.0' \
@@ -254,7 +254,7 @@ steps:
254 - vllm/model_executor/guided_decoding
255 - tests/test_logits_processor
256 - tests/model_executor/test_guided_processors
257 - commands:
257 + commands:
258 - pytest -v -s test_logits_processor.py
259 - pytest -v -s model_executor/test_guided_processors.py
260
@@ -265,7 +265,7 @@ steps:
265 - vllm/model_executor/models/eagle.py
266 commands:
267 - pytest -v -s spec_decode/e2e/test_multistep_correctness.py
268 - - VLLM_ATTENTION_BACKEND=FLASH_ATTN pytest -v -s spec_decode --ignore=spec_decode/e2e/test_multistep_correctness.py
268 + - VLLM_ATTENTION_BACKEND=FLASH_ATTN pytest -v -s spec_decode --ignore=spec_decode/e2e/test_multistep_correctness.py --ignore=spec_decode/e2e/test_mtp_correctness.py
269 - pytest -v -s spec_decode/e2e/test_eagle_correctness.py
270
271 - label: LoRA Test %N # 15min each
@@ -580,7 +580,7 @@ steps:
580 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
581 # This test runs llama 13B, so it is required to run on 4 GPUs.
582 - pytest -v -s -x lora/test_long_context.py
583 - # There is some Tensor Parallelism related processing logic in LoRA that
583 + # There is some Tensor Parallelism related processing logic in LoRA that
584 # requires multi-GPU testing for validation.
585 - pytest -v -s -x lora/test_chatglm3_tp.py
586 - pytest -v -s -x lora/test_llama_tp.py
@@ -605,7 +605,7 @@ steps:
605 - vllm/
606 - tests/weight_loading
607 commands:
608 - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
608 + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
609
610
611 ##### multi gpus test #####
@@ -617,7 +617,7 @@ steps:
617 num_gpus: 4
618 source_file_dependencies:
619 - vllm/
620 - commands:
620 + commands:
621 # NOTE: don't test llama model here, it seems hf implementation is buggy
622 # see https://github.com/vllm-project/vllm/pull/5689 for details
623 - pytest -v -s distributed/test_custom_all_reduce.py
tests/models/registry.py
+3
@@ -296,6 +296,9 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
296 speculative_model="abhigoyal/vllm-medusa-llama-68m-random"), # noqa: E501
297 "MLPSpeculatorPreTrainedModel": _HfExamplesInfo("JackFram/llama-160m",
298 speculative_model="ibm-ai-platform/llama-160m-accelerator"), # noqa: E501
299 + "DeepSeekMTPModel": _HfExamplesInfo("luccafong/deepseek_mtp_main_random",
300 + speculative_model="luccafong/deepseek_mtp_draft_random", # noqa: E501
301 + trust_remote_code=True),
302 }
303
304 _FALLBACK_MODEL = {
tests/spec_decode/e2e/test_mtp_correctness.py new
+318
@@ -0,0 +1,318 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""This docstring details important information on the testing methodology.
3 +
4 +Most of the tests rely on "greedy equality", where we expect the output of
5 +speculative decoding on a sequence to exactly match the output of normal non-
6 +speculative decoding.
7 +
8 +Since speculative decoding with rejection sampling guarantees that the output
9 +distribution matches the target model's output distribution (up to hardware
10 +numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
11 +equality.
12 +
13 +However, we still need to verify below scenario could be passed:
14 + * Batch size 1 greedy equality
15 + * Batch size >1 greedy equality
16 + * Test greedy equality under preemption
17 + * Test greedy equality under various number of speculative tokens.
18 +
19 +With those tests, we can say at least, mtp would not break the
20 +correctess for the target model outputs.
21 +"""
22 +
23 +import pytest
24 +
25 +from .conftest import run_equality_correctness_test
26 +
27 +# main model
28 +MAIN_MODEL = "luccafong/deepseek_mtp_main_random"
29 +
30 +# max. number of speculative tokens: this corresponds to
31 +# num_nextn_predict_layers in the config.json of the speculator model.
32 +MAX_SPEC_TOKENS = 1
33 +
34 +# precision
35 +PRECISION = "bfloat16"
36 +
37 +
38 +@pytest.mark.parametrize(
39 + "common_llm_kwargs",
40 + [{
41 + # Skip cuda graph recording for fast test.
42 + "enforce_eager": True,
43 +
44 + # Print spec metrics.
45 + "disable_log_stats": False,
46 +
47 + # Precision
48 + "dtype": PRECISION,
49 +
50 + # Main model
51 + "model_name": MAIN_MODEL,
52 +
53 + # GPU memory utilization
54 + "gpu_memory_utilization": 0.85
55 + }])
56 +@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
57 +@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
58 +@pytest.mark.parametrize("test_llm_kwargs", [
59 + {
60 + "num_speculative_tokens": MAX_SPEC_TOKENS,
61 + },
62 +])
63 +@pytest.mark.parametrize("output_len", [
64 + 128,
65 +])
66 +@pytest.mark.parametrize("batch_size", [1, 32])
67 +@pytest.mark.parametrize("seed", [1])
68 +def test_mtp_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
69 + per_test_common_llm_kwargs,
70 + baseline_llm_kwargs, test_llm_kwargs,
71 + batch_size: int, output_len: int,
72 + seed: int):
73 +
74 + run_equality_correctness_test(vllm_runner, common_llm_kwargs,
75 + per_test_common_llm_kwargs,
76 + baseline_llm_kwargs, test_llm_kwargs,
77 + batch_size, output_len, seed)
78 +
79 +
80 +@pytest.mark.parametrize(
81 + "common_llm_kwargs",
82 + [{
83 + # Skip cuda graph recording for fast test.
84 + "enforce_eager": True,
85 +
86 + # Print spec metrics.
87 + "disable_log_stats": False,
88 +
89 + # Precision
90 + "dtype": PRECISION,
91 +
92 + # Main model
93 + "model_name": MAIN_MODEL,
94 +
95 + # GPU memory utilization
96 + "gpu_memory_utilization": 0.85
97 + }])
98 +@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
99 +@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
100 +@pytest.mark.parametrize("test_llm_kwargs", [
101 + {
102 + "num_speculative_tokens": MAX_SPEC_TOKENS,
103 + "disable_logprobs_during_spec_decoding": False,
104 + },
105 + {
106 + "num_speculative_tokens": MAX_SPEC_TOKENS,
107 + "disable_logprobs_during_spec_decoding": True,
108 + },
109 +])
110 +@pytest.mark.parametrize("output_len", [
111 + 128,
112 +])
113 +@pytest.mark.parametrize("batch_size", [8])
114 +@pytest.mark.parametrize("seed", [1])
115 +@pytest.mark.parametrize("logprobs", [1, 6])
116 +def test_mtp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
117 + per_test_common_llm_kwargs,
118 + baseline_llm_kwargs, test_llm_kwargs,
119 + batch_size: int, output_len: int, seed: int,
120 + logprobs: int):
121 +
122 + run_equality_correctness_test(vllm_runner,
123 + common_llm_kwargs,
124 + per_test_common_llm_kwargs,
125 + baseline_llm_kwargs,
126 + test_llm_kwargs,
127 + batch_size,
128 + output_len,
129 + seed,
130 + logprobs=logprobs,
131 + prompt_logprobs=logprobs,
132 + disable_logprobs=test_llm_kwargs[
133 + 'disable_logprobs_during_spec_decoding'])
134 +
135 +
136 +@pytest.mark.parametrize(
137 + "common_llm_kwargs",
138 + [{
139 + "enforce_eager": False,
140 +
141 + # Print spec metrics.
142 + "disable_log_stats": False,
143 +
144 + # Precision
145 + "dtype": PRECISION,
146 +
147 + # Main model
148 + "model_name": MAIN_MODEL,
149 + "gpu_memory_utilization": 0.85
150 + }])
151 +@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
152 +@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
153 +@pytest.mark.parametrize("test_llm_kwargs", [
154 + {
155 + "num_speculative_tokens": MAX_SPEC_TOKENS,
156 + },
157 +])
158 +@pytest.mark.parametrize("output_len", [
159 + 128,
160 +])
161 +@pytest.mark.parametrize("batch_size", [1, 32])
162 +@pytest.mark.parametrize("seed", [1])
163 +def test_mtp_e2e_greedy_correctness_cuda_graph(vllm_runner, common_llm_kwargs,
164 + per_test_common_llm_kwargs,
165 + baseline_llm_kwargs,
166 + test_llm_kwargs,
167 + batch_size: int,
168 + output_len: int, seed: int):
169 + """Verify greedy equality with cuda graph enabled and different
170 + batch sizes."""
171 + run_equality_correctness_test(vllm_runner, common_llm_kwargs,
172 + per_test_common_llm_kwargs,
173 + baseline_llm_kwargs, test_llm_kwargs,
174 + batch_size, output_len, seed)
175 +
176 +
177 +@pytest.mark.parametrize(
178 + "common_llm_kwargs",
179 + [{
180 + "block_size": 8,
181 + # 2 for small prompt, 256//8 for generated.
182 + "num_gpu_blocks_override": 2 + 256 // 8,
183 + "max_model_len": (2 + 256 // 8) * 8,
184 +
185 + # Skip cuda graph recording for fast test.
186 + "enforce_eager": True,
187 +
188 + # Precision
189 + "dtype": PRECISION,
190 +
191 + # Main model
192 + "model_name": MAIN_MODEL,
193 +
194 + # GPU memory utilization
195 + "gpu_memory_utilization": 0.9
196 + }])
197 +@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
198 +@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
199 +@pytest.mark.parametrize("test_llm_kwargs", [
200 + {
201 + "num_speculative_tokens": MAX_SPEC_TOKENS,
202 + },
203 +])
204 +@pytest.mark.parametrize(
205 + "output_len",
206 + [
207 + # Use small output len for fast test.
208 + 128,
209 + ])
210 +@pytest.mark.parametrize("batch_size", [4])
211 +@pytest.mark.parametrize("seed", [1])
212 +def test_mtp_e2e_greedy_correctness_with_preemption(
213 + vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
214 + baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
215 + seed: int):
216 + """Verify greedy equality, even when some sequences are preempted mid-
217 + generation.
218 + """
219 + run_equality_correctness_test(vllm_runner, common_llm_kwargs,
220 + per_test_common_llm_kwargs,
221 + baseline_llm_kwargs, test_llm_kwargs,
222 + batch_size, output_len, seed)
223 +
224 +
225 +@pytest.mark.parametrize(
226 + "common_llm_kwargs",
227 + [{
228 + # Skip cuda graph recording for fast test.
229 + "enforce_eager": True,
230 +
231 + # Precision
232 + "dtype": PRECISION,
233 +
234 + # Main model
235 + "model_name": MAIN_MODEL,
236 +
237 + # GPU memory utilization
238 + "gpu_memory_utilization": 0.9
239 + }])
240 +@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
241 +@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
242 +@pytest.mark.parametrize(
243 + "test_llm_kwargs",
244 + [
245 + {
246 + "num_speculative_tokens": k,
247 + }
248 + # Try a range of num. speculative tokens
249 + for k in range(1, 1 + MAX_SPEC_TOKENS)
250 + ])
251 +@pytest.mark.parametrize("batch_size", [2])
252 +@pytest.mark.parametrize(
253 + "output_len",
254 + [
255 + # Use smaller output len for fast test.
256 + 32,
257 + ])
258 +@pytest.mark.parametrize("seed", [1])
259 +def test_mtp_different_k(vllm_runner, common_llm_kwargs,
260 + per_test_common_llm_kwargs, baseline_llm_kwargs,
261 + test_llm_kwargs, batch_size: int, output_len: int,
262 + seed: int):
263 + """Verify that mtp speculative decoding produces exact equality
264 + to without spec decode with different values of num_speculative_tokens.
265 + """
266 + run_equality_correctness_test(vllm_runner, common_llm_kwargs,
267 + per_test_common_llm_kwargs,
268 + baseline_llm_kwargs, test_llm_kwargs,
269 + batch_size, output_len, seed)
270 +
271 +
272 +@pytest.mark.parametrize(
273 + "common_llm_kwargs",
274 + [{
275 + # Skip cuda graph recording for fast test.
276 + "enforce_eager": True,
277 +
278 + # Precision
279 + "dtype": PRECISION,
280 +
281 + # Main model
282 + "model_name": MAIN_MODEL,
283 +
284 + # GPU memory utilization
285 + "gpu_memory_utilization": 0.9
286 + }])
287 +@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
288 +@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
289 +@pytest.mark.parametrize("test_llm_kwargs",
290 + [{
291 + "num_speculative_tokens": MAX_SPEC_TOKENS,
292 + "speculative_disable_by_batch_size": 4
293 + }])
294 +@pytest.mark.parametrize("batch_size", [1, 5])
295 +@pytest.mark.parametrize(
296 + "output_len",
297 + [
298 + # Use smaller output len for fast test.
299 + 32,
300 + ])
301 +@pytest.mark.parametrize("seed", [1])
302 +def test_mtp_disable_queue(vllm_runner, common_llm_kwargs,
303 + per_test_common_llm_kwargs, baseline_llm_kwargs,
304 + test_llm_kwargs, batch_size: int, output_len: int,
305 + seed: int):
306 + """Verify that mtp speculative decoding produces exact equality
307 + to without spec decode when speculation is disabled for large
308 + batch sizes.
309 + """
310 + run_equality_correctness_test(vllm_runner, common_llm_kwargs,
311 + per_test_common_llm_kwargs,
312 + baseline_llm_kwargs, test_llm_kwargs,
313 + batch_size, output_len, seed)
314 +
315 +
316 +if __name__ == "__main__":
317 + import pytest
318 + pytest.main([__file__])
vllm/config.py
+33 -10
@@ -763,7 +763,7 @@ class ModelConfig:
763 def is_deepseek_mla(self) -> bool:
764 return (hasattr(self.hf_text_config, "model_type")) \
765 and (self.hf_text_config.model_type in \
766 - ('deepseek_v2', 'deepseek_v3'))\
766 + ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'))\
767 and (self.hf_text_config.kv_lora_rank is not None)
768
769 def get_head_size(self) -> int:
@@ -856,8 +856,12 @@ class ModelConfig:
856 def get_layers_start_end_indices(
857 self, parallel_config: "ParallelConfig") -> Tuple[int, int]:
858 from vllm.distributed.utils import get_pp_indices
859 - total_num_hidden_layers = getattr(self.hf_text_config,
860 - "num_hidden_layers", 0)
859 + if self.hf_text_config.model_type == "deepseek_mtp":
860 + total_num_hidden_layers = getattr(self.hf_text_config,
861 + "num_nextn_predict_layers", 0)
862 + else:
863 + total_num_hidden_layers = getattr(self.hf_text_config,
864 + "num_hidden_layers", 0)
865 pp_rank = parallel_config.rank // parallel_config.tensor_parallel_size
866 pp_size = parallel_config.pipeline_parallel_size
867 start, end = get_pp_indices(total_num_hidden_layers, pp_rank, pp_size)
@@ -1689,6 +1693,18 @@ class SpeculativeConfig:
1693 hash_str = hashlib.md5(str(factors).encode()).hexdigest()
1694 return hash_str
1695
1696 + @staticmethod
1697 + def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
1698 + if hf_config.model_type == "deepseek_v3":
1699 + hf_config.model_type = "deepseek_mtp"
1700 + if hf_config.model_type == "deepseek_mtp":
1701 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
1702 + hf_config.update({
1703 + "n_predict": n_predict,
1704 + "architectures": ["DeepSeekMTPModel"]
1705 + })
1706 + return hf_config
1707 +
1708 @staticmethod
1709 def maybe_create_spec_config(
1710 target_model_config: ModelConfig,
@@ -1771,12 +1787,18 @@ class SpeculativeConfig:
1787 Optional["SpeculativeConfig"]: An instance of SpeculativeConfig if
1788 the necessary conditions are met, else None.
1789 """
1774 -
1790 if speculative_model is None:
1791 if num_speculative_tokens is not None:
1777 - raise ValueError("num_speculative_tokens was provided without "
1778 - "speculative_model.")
1779 - return None
1792 + if target_model_config.hf_text_config.model_type \
1793 + == "deepseek_v3":
1794 + # use the draft model from the same model:
1795 + speculative_model = target_model_config.model
1796 + else:
1797 + raise ValueError(
1798 + "num_speculative_tokens was provided without "
1799 + "speculative_model.")
1800 + else:
1801 + return None
1802
1803 if (speculative_disable_by_batch_size is not None
1804 and speculative_disable_by_batch_size < 2):
@@ -1830,6 +1852,7 @@ class SpeculativeConfig:
1852 max_seq_len_to_capture=target_model_config.
1853 max_seq_len_to_capture,
1854 max_logprobs=target_model_config.max_logprobs,
1855 + hf_overrides=SpeculativeConfig.hf_config_override,
1856 )
1857
1858 draft_hf_config = draft_model_config.hf_config
@@ -1846,7 +1869,6 @@ class SpeculativeConfig:
1869 if (num_speculative_tokens is not None
1870 and hasattr(draft_hf_config, "num_lookahead_tokens")):
1871 draft_hf_config.num_lookahead_tokens = num_speculative_tokens
1849 -
1872 n_predict = getattr(draft_hf_config, "n_predict", None)
1873 if n_predict is not None:
1874 if num_speculative_tokens is None:
@@ -1960,8 +1982,9 @@ class SpeculativeConfig:
1982 speculative_draft_tensor_parallel_size = 1
1983 if target_parallel_config.tensor_parallel_size > 1:
1984 logger.warning(
1963 - "MLPSpeculator cannot currently be run with tp>1; "
1964 - "setting speculative_draft_tensor_parallel_size=1")
1985 + "%s cannot currently be run with tp>1; "
1986 + "setting speculative_draft_tensor_parallel_size=1",
1987 + draft_hf_config.model_type)
1988 else:
1989 speculative_draft_tensor_parallel_size = \
1990 target_parallel_config.tensor_parallel_size
vllm/model_executor/models/deepseek_mtp.py new
+284
@@ -0,0 +1,284 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Iterable, List, Optional, Set, Tuple
3 +
4 +import torch
5 +import torch.nn as nn
6 +from transformers import PretrainedConfig
7 +
8 +from vllm.attention.backends.abstract import AttentionMetadata
9 +from vllm.config import CacheConfig, ModelConfig, VllmConfig
10 +from vllm.model_executor.layers.fused_moe import FusedMoE
11 +from vllm.model_executor.layers.layernorm import RMSNorm
12 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
13 +from vllm.model_executor.layers.quantization import QuantizationConfig
14 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
15 +from vllm.model_executor.layers.vocab_parallel_embedding import (
16 + ParallelLMHead, VocabParallelEmbedding)
17 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
18 +from vllm.model_executor.sampling_metadata import SamplingMetadata
19 +from vllm.sequence import IntermediateTensors
20 +
21 +from .deepseek_v2 import (DeepseekV2DecoderLayer,
22 + get_spec_layer_idx_from_weight_name)
23 +from .utils import maybe_prefix
24 +
25 +
26 +class SharedHead(nn.Module):
27 +
28 + def __init__(
29 + self,
30 + config: PretrainedConfig,
31 + quant_config: Optional[QuantizationConfig] = None,
32 + ) -> None:
33 + super().__init__()
34 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
35 + self.head = ParallelLMHead(config.vocab_size,
36 + config.hidden_size,
37 + quant_config=quant_config)
38 +
39 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
40 + return self.norm(hidden_states)
41 +
42 +
43 +class DeepSeekMultiTokenPredictorLayer(nn.Module):
44 +
45 + def __init__(
46 + self,
47 + config: PretrainedConfig,
48 + prefix: str,
49 + model_config: ModelConfig,
50 + cache_config: Optional[CacheConfig] = None,
51 + quant_config: Optional[QuantizationConfig] = None,
52 + ) -> None:
53 + super().__init__()
54 + self.embed_tokens = VocabParallelEmbedding(
55 + config.vocab_size,
56 + config.hidden_size,
57 + )
58 +
59 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
60 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
61 + self.eh_proj = nn.Linear(config.hidden_size * 2,
62 + config.hidden_size,
63 + bias=False)
64 + self.shared_head = SharedHead(config=config, quant_config=quant_config)
65 + self.mtp_block = DeepseekV2DecoderLayer(config, prefix, model_config,
66 + cache_config, quant_config)
67 +
68 + def forward(
69 + self,
70 + input_ids: torch.Tensor,
71 + positions: torch.Tensor,
72 + kv_cache: torch.Tensor,
73 + attn_metadata: AttentionMetadata,
74 + previous_hidden_states: torch.Tensor,
75 + inputs_embeds: Optional[torch.Tensor] = None,
76 + spec_step_index: int = 0,
77 + ) -> torch.Tensor:
78 + if inputs_embeds is None:
79 + inputs_embeds = self.embed_tokens(input_ids)
80 + assert inputs_embeds is not None
81 + # masking inputs at position 0, as not needed by MTP
82 + inputs_embeds[positions == 0] = 0
83 + inputs_embeds = self.enorm(inputs_embeds)
84 + previous_hidden_states = self.hnorm(previous_hidden_states)
85 +
86 + hidden_states = self.eh_proj(
87 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
88 +
89 + hidden_states, residual = self.mtp_block(positions=positions,
90 + hidden_states=hidden_states,
91 + kv_cache=kv_cache,
92 + attn_metadata=attn_metadata,
93 + residual=None)
94 + hidden_states = residual + hidden_states
95 + return self.shared_head(hidden_states)
96 +
97 +
98 +class DeepSeekMultiTokenPredictor(nn.Module):
99 +
100 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
101 + super().__init__()
102 + config = vllm_config.model_config.hf_config
103 + self.mtp_start_layer_idx = config.num_hidden_layers
104 + self.num_mtp_layers = config.num_nextn_predict_layers
105 + # to map the exact layer index from weights
106 + self.layers = torch.nn.ModuleDict({
107 + str(idx):
108 + DeepSeekMultiTokenPredictorLayer(
109 + config,
110 + f"{prefix}.layers.{idx}",
111 + model_config=vllm_config.model_config,
112 + cache_config=vllm_config.cache_config,
113 + quant_config=vllm_config.quant_config,
114 + )
115 + for idx in range(self.mtp_start_layer_idx,
116 + self.mtp_start_layer_idx + self.num_mtp_layers)
117 + })
118 +
119 + self.logits_processor = LogitsProcessor(config.vocab_size)
120 +
121 + def forward(
122 + self,
123 + input_ids: torch.Tensor,
124 + positions: torch.Tensor,
125 + kv_caches: List[torch.Tensor],
126 + attn_metadata: AttentionMetadata,
127 + previous_hidden_states: torch.Tensor,
128 + inputs_embeds: Optional[torch.Tensor] = None,
129 + spec_step_idx: int = 0,
130 + ) -> torch.Tensor:
131 + return self.layers[str(self.mtp_start_layer_idx + spec_step_idx)](
132 + input_ids,
133 + positions,
134 + kv_caches[spec_step_idx],
135 + attn_metadata,
136 + previous_hidden_states,
137 + inputs_embeds,
138 + spec_step_idx,
139 + )
140 +
141 + def compute_logits(
142 + self,
143 + hidden_states: torch.Tensor,
144 + sampling_metadata: SamplingMetadata,
145 + spec_step_idx: int = 0,
146 + ) -> torch.Tensor:
147 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + spec_step_idx)]
148 + logits = self.logits_processor(mtp_layer.shared_head.head,
149 + hidden_states, sampling_metadata)
150 + return logits
151 +
152 +
153 +class DeepSeekMTP(nn.Module):
154 +
155 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
156 + super().__init__()
157 + self.config = vllm_config.model_config.hf_config
158 + self.model = DeepSeekMultiTokenPredictor(vllm_config=vllm_config,
159 + prefix=maybe_prefix(
160 + prefix, "model"))
161 +
162 + self.sampler = get_sampler()
163 +
164 + def forward(
165 + self,
166 + input_ids: torch.Tensor,
167 + positions: torch.Tensor,
168 + kv_caches: List[torch.Tensor],
169 + attn_metadata: AttentionMetadata,
170 + previous_hidden_states: torch.Tensor,
171 + intermediate_tensors: Optional[IntermediateTensors] = None,
172 + inputs_embeds: Optional[torch.Tensor] = None,
173 + spec_step_idx: int = 0,
174 + ) -> torch.Tensor:
175 + hidden_states = self.model(input_ids, positions, kv_caches,
176 + attn_metadata, previous_hidden_states,
177 + inputs_embeds, spec_step_idx)
178 + return hidden_states
179 +
180 + def compute_logits(
181 + self,
182 + hidden_states: torch.Tensor,
183 + sampling_metadata: SamplingMetadata,
184 + spec_step_idx: int = 0,
185 + ) -> Optional[torch.Tensor]:
186 + return self.model.compute_logits(hidden_states, sampling_metadata,
187 + spec_step_idx)
188 +
189 + def sample(
190 + self,
191 + logits: torch.Tensor,
192 + sampling_metadata: SamplingMetadata,
193 + ) -> Optional[SamplerOutput]:
194 + next_tokens = self.sampler(logits, sampling_metadata)
195 + return next_tokens
196 +
197 + def load_weights(self, weights: Iterable[Tuple[str,
198 + torch.Tensor]]) -> Set[str]:
199 + stacked_params_mapping = [
200 + ("gate_up_proj", "gate_proj", 0),
201 + ("gate_up_proj", "up_proj", 1),
202 + ]
203 +
204 + expert_params_mapping = FusedMoE.make_expert_params_mapping(
205 + ckpt_gate_proj_name="gate_proj",
206 + ckpt_down_proj_name="down_proj",
207 + ckpt_up_proj_name="up_proj",
208 + num_experts=self.config.n_routed_experts)
209 +
210 + params_dict = dict(self.named_parameters())
211 + loaded_params: Set[str] = set()
212 + for name, loaded_weight in weights:
213 + if "rotary_emb.inv_freq" in name:
214 + continue
215 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
216 + if spec_layer is None:
217 + continue
218 + name = self._rewrite_spec_layer_name(spec_layer, name)
219 + for (param_name, weight_name, shard_id) in stacked_params_mapping:
220 + # Skip non-stacked layers and experts (experts handled below).
221 + if weight_name not in name:
222 + continue
223 + # We have mlp.experts[0].gate_proj in the checkpoint.
224 + # Since we handle the experts below in expert_params_mapping,
225 + # we need to skip here BEFORE we update the name, otherwise
226 + # name will be updated to mlp.experts[0].gate_up_proj, which
227 + # will then be updated below in expert_params_mapping
228 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
229 + if (("mlp.experts." in name) and name not in params_dict):
230 + continue
231 + name = name.replace(weight_name, param_name)
232 + # Skip loading extra bias for GPTQ models.
233 + if name.endswith(".bias") and name not in params_dict:
234 + continue
235 +
236 + param = params_dict[name]
237 + weight_loader = param.weight_loader
238 + weight_loader(param, loaded_weight, shard_id)
239 + break
240 + else:
241 + for mapping in expert_params_mapping:
242 + param_name, weight_name, expert_id, shard_id = mapping
243 + if weight_name not in name:
244 + continue
245 + name = name.replace(weight_name, param_name)
246 +
247 + param = params_dict[name]
248 + weight_loader = param.weight_loader
249 + weight_loader(param,
250 + loaded_weight,
251 + name,
252 + shard_id=shard_id,
253 + expert_id=expert_id)
254 + break
255 + else:
256 + # Skip loading extra bias for GPTQ models.
257 + if name.endswith(".bias") and name not in params_dict:
258 + continue
259 +
260 + param = params_dict[name]
261 + weight_loader = getattr(param, "weight_loader",
262 + default_weight_loader)
263 + weight_loader(param, loaded_weight)
264 + loaded_params.add(name)
265 + return loaded_params
266 +
267 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
268 + """
269 + Rewrite the weight name to match the format of the original model.
270 + Add .mtp_block for modules in transformer layer block for spec layer
271 + """
272 + spec_layer_weight_names = [
273 + "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"
274 + ]
275 + spec_layer_weight = False
276 + for weight_name in spec_layer_weight_names:
277 + if weight_name in name:
278 + spec_layer_weight = True
279 + break
280 + if not spec_layer_weight:
281 + # treat rest weights as weights for transformer layer block
282 + name = name.replace(f"model.layers.{spec_layer}.",
283 + f"model.layers.{spec_layer}.mtp_block.")
284 + return name
vllm/model_executor/models/deepseek_v2.py
+15 -7
@@ -732,13 +732,9 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
732 if "rotary_emb.inv_freq" in name:
733 continue
734
735 - # TODO(simon): support nextn predict layers
736 - if hasattr(self.config, "num_nextn_predict_layers"
737 - ) and self.config.num_nextn_predict_layers > 0:
738 - assert self.config.num_nextn_predict_layers == 1
739 - layer_idx = self.config.num_hidden_layers
740 - if name.startswith(f"model.layers.{layer_idx}"):
741 - continue
735 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
736 + if spec_layer is not None:
737 + continue # skip spec decode layers for main model
738
739 for (param_name, weight_name, shard_id) in stacked_params_mapping:
740 # Skip non-stacked layers and experts (experts handled below).
@@ -805,3 +801,15 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
801
802 class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
803 pass
804 +
805 +
806 +def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,
807 + weight_name: str) -> Optional[int]:
808 + if hasattr(config,
809 + "num_nextn_predict_layers") and (config.num_nextn_predict_layers
810 + > 0):
811 + layer_idx = config.num_hidden_layers
812 + for i in range(config.num_nextn_predict_layers):
813 + if weight_name.startswith(f"model.layers.{layer_idx+i}."):
814 + return layer_idx + i
815 + return None
vllm/model_executor/models/registry.py
+1
@@ -187,6 +187,7 @@ _MULTIMODAL_MODELS = {
187
188 _SPECULATIVE_DECODING_MODELS = {
189 "EAGLEModel": ("eagle", "EAGLE"),
190 + "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
191 "MedusaModel": ("medusa", "Medusa"),
192 "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
193 }
vllm/sequence.py
+2
@@ -1307,6 +1307,8 @@ class ExecuteModelRequest(
1307 previous_hidden_states: Optional[HiddenStates] = None
1308 # The number of forward steps to run.
1309 num_steps: int = 1
1310 + # The step index for spec model input.
1311 + spec_step_idx: Optional[int] = None
1312 # Finished request ids since last step.
1313 finished_requests_ids: List[str] = msgspec.field(default_factory=list)
1314 # The last sampled token ids for multi step decoding.
vllm/spec_decode/draft_model_runner.py
+15 -5
@@ -153,7 +153,7 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
153 return False
154
155 # TODO: Add support for other attn backends
156 - if self.attn_backend.get_name() != "FLASH_ATTN":
156 + if self.attn_backend.get_name() not in ("FLASH_ATTN", "TRITON_MLA"):
157 return False
158
159 # TODO: Add support for LORA
@@ -175,6 +175,7 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
175 previous_hidden_states: Optional[torch.Tensor] = None,
176 intermediate_tensors: Optional[IntermediateTensors] = None,
177 num_steps: int = 1,
178 + **kwargs,
179 ) -> Optional[List[SamplerOutput]]:
180 """Executes num_steps forward passes with advacement of input tensors
181 on the GPU. Look at supports_gpu_multi_step(..) for pre-conditions.
@@ -271,10 +272,17 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
272 for step in range(num_steps):
273 multi_modal_kwargs = model_input.multi_modal_kwargs or {}
274
274 - kwargs = {"previous_hidden_states": hidden_states} \
275 + model_execute_kwargs = {"previous_hidden_states": hidden_states} \
276 if previous_hidden_states is not None else {}
277
278 + compute_logits_kwargs = {}
279 # Run model
280 + if hasattr(self.model.config, "num_nextn_predict_layers"):
281 + # for DeepSeek MTP only to use the corresponding layer for
282 + # each step
283 + spec_step_idx = kwargs.get("spec_step_idx", step)
284 + model_execute_kwargs["spec_step_idx"] = spec_step_idx
285 + compute_logits_kwargs["spec_step_idx"] = spec_step_idx
286 with set_forward_context(model_input.attn_metadata,
287 self.vllm_config):
288 hidden_states = model_executable(
@@ -285,13 +293,15 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
293 intermediate_tensors=intermediate_tensors,
294 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
295 device=self.device),
288 - **kwargs,
296 + **model_execute_kwargs,
297 )
298
299 # Compute the logits.
300 logits = self.model.compute_logits(hidden_states,
293 - model_input.sampling_metadata)
294 -
301 + model_input.sampling_metadata,
302 + **compute_logits_kwargs)
303 + if not self.is_driver_worker:
304 + return []
305 # Sample the next token.
306 output = self.model.sample(
307 logits=logits,
vllm/spec_decode/spec_decode_worker.py
+20 -5
@@ -108,6 +108,7 @@ def create_spec_worker(*args, **kwargs) -> "SpecDecodeWorker":
108 typical_acceptance_sampler_posterior_alpha,
109 disable_logprobs=speculative_config.disable_logprobs,
110 disable_log_stats=speculative_config.disable_log_stats,
111 + num_speculative_tokens=speculative_config.num_speculative_tokens,
112 )
113
114 return spec_decode_worker
@@ -153,10 +154,12 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
154 typical_acceptance_sampler_posterior_alpha: float,
155 disable_logprobs: bool,
156 disable_log_stats: bool,
157 + num_speculative_tokens: int,
158 ) -> "SpecDecodeWorker":
159
160 allow_zero_draft_token_step = True
161 enable_lm_head_weight_load = False
162 + num_spec_prefill_steps = 1
163 ngram_prompt_lookup_max = (
164 draft_worker_kwargs.pop("ngram_prompt_lookup_max"))
165 ngram_prompt_lookup_min = (
@@ -179,14 +182,16 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
182 elif draft_model_config.hf_config.model_type == "medusa":
183 proposer_worker = MedusaWorker(**draft_worker_kwargs)
184 else:
182 - if draft_tp == 1:
185 + if draft_tp == 1 or draft_model_config.hf_config.model_type ==\
186 + "deepseek_mtp":
187 if current_platform.is_cuda_alike():
188 draft_worker_kwargs[
189 "model_runner_cls"] = TP1DraftModelRunner
190 else:
191 if draft_model_config.hf_config.model_type == "eagle":
192 raise NotImplementedError(
189 - "EAGLE does not support TP > 1 yet")
193 + f"{draft_model_config.hf_config.model_type} "
194 + "does not support TP > 1 yet")
195
196 allow_zero_draft_token_step = False
197
@@ -195,6 +200,8 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
200 enable_lm_head_weight_load = True
201
202 proposer_worker = MultiStepWorker(**draft_worker_kwargs)
203 + if draft_model_config.hf_config.model_type == "deepseek_mtp":
204 + num_spec_prefill_steps = num_speculative_tokens
205
206 proposer_worker = SmallerTpProposerWorker.maybe_wrap_worker(
207 proposer_worker, draft_tp, target_tp)
@@ -247,7 +254,8 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
254 disable_by_batch_size=disable_by_batch_size,
255 spec_decode_sampler=spec_decode_sampler,
256 allow_zero_draft_token_step=allow_zero_draft_token_step,
250 - enable_lm_head_weight_load=enable_lm_head_weight_load)
257 + enable_lm_head_weight_load=enable_lm_head_weight_load,
258 + num_spec_prefill_steps=num_spec_prefill_steps)
259
260 def __init__(
261 self,
@@ -261,6 +269,7 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
269 disable_by_batch_size: Optional[int] = None,
270 allow_zero_draft_token_step: Optional[bool] = True,
271 enable_lm_head_weight_load: Optional[bool] = False,
272 + num_spec_prefill_steps: int = 1,
273 ):
274 """
275 Create a SpecDecodeWorker.
@@ -293,6 +302,10 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
302 draft model is larger than 1 (TODO: #5814)
303 enable_lm_head_weight_load: whether to load lm_head weight for
304 draft models like eagle.
305 + num_spec_prefill_steps: number of speculative prefill steps to run
306 + before the speculative decoding starts. This is only used when
307 + the draft model is a deepseek_mtp model that requires prefill
308 + kv cache separately for each MTP layer.
309 """
310 self.proposer_worker = proposer_worker
311 self.scorer_worker = scorer_worker
@@ -326,6 +339,7 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
339 self.previous_hidden_states: Optional[HiddenStates] = None
340 self._disable_logprobs = disable_logprobs
341 self._disable_log_stats = disable_log_stats
342 + self._num_spec_prefill_steps = num_spec_prefill_steps
343
344 def init_device(self) -> None:
345 """Initialize both scorer and proposer models.
@@ -685,8 +699,9 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
699 execute_model_req.previous_hidden_states = \
700 prepare_prefill_hidden_states(
701 sampler_output.prefill_hidden_states)
688 -
689 - self.proposer_worker.execute_model(execute_model_req)
702 + for i in range(self._num_spec_prefill_steps):
703 + execute_model_req.spec_step_idx = i
704 + self.proposer_worker.execute_model(execute_model_req)
705
706 sampler_output_to_return = (self._serialize_sampler_output_no_logprobs(
707 execute_model_req=execute_model_req, sampler_output=sampler_output)
vllm/worker/model_runner.py
+16 -4
@@ -99,6 +99,7 @@ class ModelInputForGPU(ModelRunnerInputBase):
99 virtual_engine: int = 0
100 async_callback: Optional[Callable] = None
101 scheduler_outputs: Optional[SchedulerOutputs] = None
102 + previous_hidden_states: Optional[torch.Tensor] = None
103
104 def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
105 tensor_dict = {
@@ -1649,6 +1650,7 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1650 kv_caches: List[torch.Tensor],
1651 intermediate_tensors: Optional[IntermediateTensors] = None,
1652 num_steps: int = 1,
1653 + **kwargs,
1654 ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:
1655 if num_steps > 1:
1656 raise ValueError("num_steps > 1 is not supported in ModelRunner")
@@ -1706,6 +1708,10 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1708 "finished_requests_ids": model_input.finished_requests_ids,
1709 "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
1710 } if self.has_inner_state else {}
1711 + previous_hidden_states = kwargs.get("previous_hidden_states")
1712 + model_kwargs = {}
1713 + if previous_hidden_states is not None:
1714 + model_kwargs["previous_hidden_states"] = previous_hidden_states
1715 if (self.observability_config is not None
1716 and self.observability_config.collect_model_forward_time):
1717 model_forward_start = torch.cuda.Event(enable_timing=True)
@@ -1723,7 +1729,9 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1729 intermediate_tensors=intermediate_tensors,
1730 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
1731 device=self.device),
1726 - **seqlen_agnostic_kwargs)
1732 + **seqlen_agnostic_kwargs,
1733 + **model_kwargs,
1734 + )
1735
1736 if (self.observability_config is not None
1737 and self.observability_config.collect_model_forward_time):
@@ -1815,7 +1823,7 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1823 1. current vLLM instance is KV cache consumer/decode vLLM instance
1824 2. this batch is not a profiling run
1825 3. this batch is a prefill run
1818 -
1826 +
1827 Args:
1828 model_input: input to the model executable
1829 kv_caches: vLLM's paged memory
@@ -1840,7 +1848,7 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1848 1. current vLLM instance is KV cache producer/prefill vLLM instance
1849 2. this batch is not a profiling run
1850 3. this batch is a prefill run
1843 -
1851 +
1852 Args:
1853 model_input: input to the model executable
1854 kv_caches: vLLM's paged memory
@@ -1976,7 +1984,11 @@ class CUDAGraphRunner(nn.Module):
1984 # Copy the input tensors to the input buffers.
1985 self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)
1986 if positions is not None:
1979 - self.input_buffers["positions"].copy_(positions, non_blocking=True)
1987 + # in some case like MLA, it will reuse positions in metadata
1988 + # but truncate them to the original size
1989 + # so the shape is not padded, we need to copy partial only
1990 + self.input_buffers["positions"][:positions.shape[0]].copy_(
1991 + positions, non_blocking=True)
1992
1993 if self.backend_name != "NO_ATTENTION":
1994 self.input_buffers["slot_mapping"].copy_(
vllm/worker/model_runner_base.py
+4 -1
@@ -46,7 +46,10 @@ def _init_attn_metadata_from_tensor_dict(
46 valid_attn_kwargs = {}
47 for field in dataclasses.fields(attn_backend.get_metadata_cls()):
48 if field.name in tensor_dict:
49 - valid_attn_kwargs[field.name] = tensor_dict.pop(field.name)
49 + if field.name == "input_positions":
50 + valid_attn_kwargs[field.name] = tensor_dict[field.name]
51 + else:
52 + valid_attn_kwargs[field.name] = tensor_dict.pop(field.name)
53
54 attn_metadata = attn_backend.make_metadata(**valid_attn_kwargs)
55 tensor_dict["attn_metadata"] = attn_metadata
vllm/worker/worker.py
+3 -3
@@ -68,10 +68,10 @@ class Worker(LocalOrDistributedWorkerBase):
68 speculative_config = self.speculative_config
69 model_config = self.model_config
70 speculative_args = {} if speculative_config is None \
71 - or (speculative_config.draft_model_config.model ==
72 - model_config.model) \
71 + or (speculative_config.draft_model_config.hf_config.model_type ==
72 + model_config.hf_config.model_type) \
73 or (speculative_config.draft_model_config.hf_config.model_type
74 - not in ["medusa", "mlp_speculator", "eagle"]) \
74 + not in ("medusa", "mlp_speculator", "eagle", "deepseek_mtp")) \
75 else {"return_hidden_states": True}
76
77 ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner
vllm/worker/worker_base.py
+2
@@ -397,6 +397,8 @@ class LocalOrDistributedWorkerBase(WorkerBase):
397
398 model_input, worker_input, kwargs = inputs
399 num_steps = worker_input.num_steps
400 + if (execute_model_req is not None and execute_model_req.spec_step_idx):
401 + kwargs["spec_step_idx"] = execute_model_req.spec_step_idx
402
403 self.execute_worker(worker_input)
404