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

A high-throughput and memory-efficient inference and serving engine for LLMs

[V1][Spec Decode] Add Dynamic SD (#32374)

Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

Ekagra Ranjan committed Jun 14, 2026 at 03:14 UTC 4ef4492e9b7a5a7ba295da783d456d45db5eb9d6
23 files changed +586 -7
docs/features/speculative_decoding/README.md
+2
@@ -17,6 +17,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
17 - [Suffix Decoding](suffix.md)
18 - [Hidden State Extraction](extract_hidden_states.md)
19 - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)
20 +- [Dynamic Speculative Decoding](dynamic_speculative_decoding.md)
21
22 ## Method Selection at a Glance
23
@@ -33,6 +34,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
34 | N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. |
35 | Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. |
36 | Custom Proposer | Varies | Varies | Bring your own proposer class (experimental). |
37 +| Dynamic Speculative Decoding | High gain | Higher than base SD method | Useful for RL or workload with fluctuating QPS |
38
39 For reproducible measurements in your environment, use
40 [`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)
docs/features/speculative_decoding/dynamic_speculative_decoding.md new
+78
@@ -0,0 +1,78 @@
1 +# Dynamic Speculative Decoding
2 +
3 +## Why is Dynamic SD needed?
4 +
5 +SD methods need to verify K tokens for each sequence during decoding. As BS increases, the effective BS becomes BS\*K which increases the compute requirement during verification. When this BS\*K goes beyond a critical BS then SD negatively impacts the decode speed (TPOT). DSD helps by tuning the K to an optimal value such that we continue to reap the benefits from SD.
6 +
7 +## Use cases
8 +
9 +* Variable concurrency workload using same deployment. K would decrease as concurrency increases.
10 +* During RL rollout where we start off with high BS but then end up with small BS due to very few long tail request which end up generating a lot of tokens stalling the progress of the current rollout. Here K would go up during the end of rollout.
11 +
12 +## `--speculative-config` schema
13 +
14 +To use Dynamic SD, add `num_speculative_tokens_per_batch_size` to the config of an SD method which is a list of list. Here, an entry is `[start_bs, end_bs, optimal_K]` which means when the concurrency is within range `[start_bs, end_bs]` then `optimal_K` number of draft tokens are used. For e.g.,
15 +
16 +```bash
17 +--speculative-config '{
18 + "method": "eagle",
19 + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
20 + "num_speculative_tokens": 3,
21 + "num_speculative_tokens_per_batch_size": [
22 + [1, 64, 3],
23 + [65, 128, 1],
24 + [129, 512, 0]
25 + ]
26 + }'
27 +```
28 +
29 +implies that:
30 +
31 +* K=3 will be used when the concurrency is in range [1, 64]
32 +* K=1 will be used when the concurrency is in range [65, 128]
33 +* K=0 will be used when the concurrency is in range [129, 512], i.e., no draft tokens will be produced.
34 +
35 +## Online Examples
36 +
37 +### Dynamic SD Eagle Drafter
38 +
39 +```bash
40 +VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \
41 + --speculative-config '{
42 + "method": "eagle",
43 + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
44 + "num_speculative_tokens": 3,
45 + "num_speculative_tokens_per_batch_size": [
46 + [1, 64, 3],
47 + [65, 128, 1],
48 + [129, 512, 0]
49 + ]
50 + }'
51 +```
52 +
53 +### Dynamic SD Eagle3 Drafter
54 +
55 +```bash
56 +VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \
57 + --speculative-config '{
58 + "method": "eagle3",
59 + "model": "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
60 + "num_speculative_tokens": 3,
61 + "num_speculative_tokens_per_batch_size": [
62 + [1, 16, 5],
63 + [17, 32, 4],
64 + [33, 64, 3],
65 + [65, 128, 1],
66 + [129, 512, 0]
67 + ]
68 + }'
69 +
70 +```
71 +
72 +## Limitations
73 +
74 +* only tested with Eagle and Eagle-3. Other SD methods may or may not work out of the box
75 +* only usable with Model Runner V1
76 +* not compatible with full cuda graph so we force piece-wise cuda graph with this feature
77 +
78 +We are working on enabling it on MRv2 with full cuda graph support.
tests/v1/spec_decode/test_dynamic_sd.py new
+218
@@ -0,0 +1,218 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Regression tests for the Dynamic SD batch-size schedule helpers."""
4 +
5 +import pytest
6 +
7 +from tests.v1.core.utils import create_requests, create_scheduler
8 +from vllm.v1.core.sched.scheduler import Scheduler
9 +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
10 +from vllm.v1.structured_output import StructuredOutputManager
11 +
12 +
13 +def _make_lookup(
14 + num_speculative_tokens_per_batch_size: list[tuple[int, int, int]],
15 + *,
16 + max_batch_size: int = 256,
17 + runtime_num_speculative_tokens: int = 3,
18 +) -> list[int]:
19 + return build_dynamic_sd_schedule_lookup(
20 + num_speculative_tokens_per_batch_size=num_speculative_tokens_per_batch_size,
21 + vllm_max_batch_size=max_batch_size,
22 + vllm_num_speculative_tokens=runtime_num_speculative_tokens,
23 + )
24 +
25 +
26 +def _make_scheduler_with_dynamic_sd(
27 + schedule: list[tuple[int, int, int]],
28 + *,
29 + max_num_seqs: int = 16,
30 + max_num_batched_tokens: int = 8192,
31 + runtime_num_speculative_tokens: int = 3,
32 +) -> Scheduler:
33 + base_scheduler = create_scheduler(
34 + max_num_seqs=max_num_seqs,
35 + max_num_batched_tokens=max_num_batched_tokens,
36 + num_speculative_tokens=runtime_num_speculative_tokens,
37 + )
38 +
39 + speculative_config = base_scheduler.vllm_config.speculative_config
40 + assert speculative_config is not None
41 + speculative_config.num_speculative_tokens_per_batch_size = schedule
42 +
43 + return Scheduler(
44 + vllm_config=base_scheduler.vllm_config,
45 + kv_cache_config=base_scheduler.kv_cache_config,
46 + block_size=base_scheduler.block_size,
47 + log_stats=True,
48 + structured_output_manager=StructuredOutputManager(base_scheduler.vllm_config),
49 + )
50 +
51 +
52 +def _add_requests_and_schedule(
53 + scheduler: Scheduler, num_requests: int, *, num_tokens: int = 10
54 +):
55 + requests = create_requests(num_requests=num_requests, num_tokens=num_tokens)
56 + for request in requests:
57 + scheduler.add_request(request)
58 + return scheduler.schedule()
59 +
60 +
61 +def test_dynamic_sd_uses_batch_size_schedule():
62 + dynamic_sd_lookup = _make_lookup(
63 + [
64 + (1, 16, 3),
65 + (32, 128, 2),
66 + (256, 2048, 0),
67 + ]
68 + )
69 +
70 + assert dynamic_sd_lookup[1] == 3
71 + assert dynamic_sd_lookup[16] == 3
72 + assert dynamic_sd_lookup[17] == 3
73 + assert dynamic_sd_lookup[31] == 3
74 + assert dynamic_sd_lookup[32] == 2
75 + assert dynamic_sd_lookup[128] == 2
76 + assert dynamic_sd_lookup[129] == 2
77 + assert dynamic_sd_lookup[255] == 2
78 + assert dynamic_sd_lookup[256] == 0
79 +
80 +
81 +def test_dynamic_sd_requires_schedule_starting_at_batch_size_one():
82 + with pytest.raises(ValueError, match="must start at 1"):
83 + _make_lookup([(2, 16, 3)])
84 +
85 +
86 +def test_dynamic_sd_clamps_k_to_runtime_max():
87 + dynamic_sd_lookup = _make_lookup(
88 + [(1, 256, 4)],
89 + runtime_num_speculative_tokens=3,
90 + )
91 +
92 + assert dynamic_sd_lookup[1] == 3
93 + assert dynamic_sd_lookup[256] == 3
94 +
95 +
96 +def test_dynamic_sd_rejects_invalid_schedule_entry():
97 + with pytest.raises(ValueError, match="3-item sequence"):
98 + _make_lookup([(1, 16, 3), (32, 64)]) # type: ignore[list-item]
99 +
100 +
101 +def test_dynamic_sd_rejects_overlapping_ranges():
102 + with pytest.raises(ValueError, match="non-overlapping and sorted"):
103 + _make_lookup([(1, 16, 3), (16, 32, 2)])
104 +
105 +
106 +def test_dynamic_sd_rejects_negative_k():
107 + with pytest.raises(ValueError, match="values must be >= 0"):
108 + _make_lookup([(1, 16, -1)])
109 +
110 +
111 +def test_dynamic_sd_rejects_empty_schedule():
112 + with pytest.raises(ValueError, match="must not be empty"):
113 + _make_lookup([])
114 +
115 +
116 +def test_dynamic_sd_requires_schedule_config():
117 + with pytest.raises(
118 + ValueError, match="num_speculative_tokens_per_batch_size is required"
119 + ):
120 + build_dynamic_sd_schedule_lookup(
121 + None,
122 + vllm_max_batch_size=256,
123 + vllm_num_speculative_tokens=3,
124 + )
125 +
126 +
127 +def test_dynamic_sd_lookup_rejects_invalid_batch_size_queries():
128 + dynamic_sd_lookup = _make_lookup([(1, 256, 3)])
129 +
130 + assert dynamic_sd_lookup[0] == 0
131 + with pytest.raises(IndexError):
132 + _ = dynamic_sd_lookup[257]
133 +
134 +
135 +def test_scheduler_initializes_dynamic_sd_lookup_from_speculative_config():
136 + scheduler = _make_scheduler_with_dynamic_sd(
137 + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)],
138 + runtime_num_speculative_tokens=3,
139 + )
140 +
141 + assert scheduler.dynamic_sd_lookup is not None
142 + assert scheduler.num_spec_tokens == 3
143 +
144 +
145 +def test_scheduler_uses_dsd_k_based_on_number_of_scheduled_requests():
146 + test_cases = [
147 + (4, 3),
148 + (64, 2),
149 + (256, 0),
150 + ]
151 +
152 + for num_requests, expected_k in test_cases:
153 + scheduler = _make_scheduler_with_dynamic_sd(
154 + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)],
155 + max_num_seqs=num_requests,
156 + max_num_batched_tokens=num_requests * 10,
157 + runtime_num_speculative_tokens=3,
158 + )
159 + output = _add_requests_and_schedule(scheduler, num_requests)
160 +
161 + assert len(output.num_scheduled_tokens) == num_requests
162 + assert output.num_spec_tokens_to_schedule == expected_k
163 +
164 +
165 +def test_scheduler_clamps_dsd_k_to_runtime_num_speculative_tokens():
166 + scheduler = _make_scheduler_with_dynamic_sd(
167 + [(1, 256, 5)],
168 + max_num_seqs=16,
169 + max_num_batched_tokens=160,
170 + runtime_num_speculative_tokens=3,
171 + )
172 + output = _add_requests_and_schedule(scheduler, 16)
173 +
174 + assert len(output.num_scheduled_tokens) == 16
175 + assert output.num_spec_tokens_to_schedule == 3
176 +
177 +
178 +def test_scheduler_falls_back_to_static_k_when_dsd_not_configured():
179 + scheduler = create_scheduler(
180 + max_num_seqs=4,
181 + max_num_batched_tokens=40,
182 + num_speculative_tokens=3,
183 + )
184 + output = _add_requests_and_schedule(scheduler, 4)
185 +
186 + assert scheduler.dynamic_sd_lookup is None
187 + assert output.num_spec_tokens_to_schedule == 3
188 +
189 +
190 +def test_scheduler_uses_static_k_when_no_requests_are_scheduled():
191 + scheduler = _make_scheduler_with_dynamic_sd(
192 + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)],
193 + runtime_num_speculative_tokens=3,
194 + )
195 + output = scheduler.schedule()
196 +
197 + assert len(output.num_scheduled_tokens) == 0
198 + assert output.num_spec_tokens_to_schedule == 3
199 +
200 +
201 +def test_scheduler_rejects_bad_dsd_config_at_construction():
202 + with pytest.raises(ValueError, match="must start at 1"):
203 + _make_scheduler_with_dynamic_sd([(2, 16, 3)])
204 +
205 +
206 +def test_scheduler_passes_max_num_seqs_as_dsd_runtime_batch_limit():
207 + scheduler = _make_scheduler_with_dynamic_sd(
208 + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)],
209 + max_num_seqs=16,
210 + max_num_batched_tokens=160,
211 + runtime_num_speculative_tokens=3,
212 + )
213 + output = _add_requests_and_schedule(scheduler, 16)
214 +
215 + assert scheduler.dynamic_sd_lookup is not None
216 + assert len(scheduler.dynamic_sd_lookup) == 17
217 + assert len(output.num_scheduled_tokens) == 16
218 + assert output.num_spec_tokens_to_schedule == 3
tests/v1/spec_decode/test_eagle.py
+2
@@ -969,6 +969,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
969 proposer.draft_attn_groups = [mock_attn_group]
970
971 result = proposer.propose(
972 + num_speculative_tokens=num_speculative_tokens,
973 target_token_ids=target_token_ids,
974 target_positions=target_positions,
975 target_hidden_states=target_hidden_states,
@@ -1071,6 +1072,7 @@ def test_propose_stores_probabilistic_draft_probs(monkeypatch):
1072 sampling_metadata.all_greedy = False
1073
1074 result = proposer.propose(
1075 + num_speculative_tokens=num_speculative_tokens,
1076 target_token_ids=torch.randint(0, vocab_size, (total_tokens,), device=device),
1077 target_positions=torch.cat(
1078 [
tests/v1/spec_decode/test_extract_hidden_states.py
+2
@@ -255,6 +255,7 @@ def test_propose():
255
256 # Call propose
257 draft_tokens = proposer.propose(
258 + num_speculative_tokens=1,
259 sampled_token_ids=sampled_token_ids,
260 target_hidden_states=target_hidden_states,
261 common_attn_metadata=common_attn_metadata,
@@ -321,6 +322,7 @@ def test_propose_different_layer_counts(num_hidden_layers):
322 ).unsqueeze(-1)
323
324 draft_tokens = proposer.propose(
325 + num_speculative_tokens=1,
326 sampled_token_ids=sampled_token_ids,
327 target_hidden_states=target_hidden_states,
328 common_attn_metadata=common_attn_metadata,
tests/v1/spec_decode/test_mtp.py
+1
@@ -205,6 +205,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):
205
206 # Run propose
207 result = proposer.propose(
208 + num_speculative_tokens=num_speculative_tokens,
209 target_token_ids=target_token_ids,
210 target_positions=target_positions,
211 target_hidden_states=target_hidden_states,
tests/v1/spec_decode/test_ngram.py
+10
@@ -81,6 +81,7 @@ def test_ngram_proposer():
81 # No match.
82 token_ids_cpu = np.array([[1, 2, 3, 4, 5]])
83 result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
84 + num_speculative_tokens=2,
85 sampled_token_ids=[[0]],
86 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
87 token_ids_cpu=token_ids_cpu,
@@ -90,6 +91,7 @@ def test_ngram_proposer():
91 # No match for 4-gram.
92 token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])
93 result = get_ngram_proposer(min_n=4, max_n=4, k=2).propose(
94 + num_speculative_tokens=2,
95 sampled_token_ids=[[0]],
96 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
97 token_ids_cpu=token_ids_cpu,
@@ -99,6 +101,7 @@ def test_ngram_proposer():
101 # No match for 4-gram but match for 3-gram.
102 token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]])
103 result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(
104 + num_speculative_tokens=2,
105 sampled_token_ids=[[0]],
106 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
107 token_ids_cpu=token_ids_cpu,
@@ -109,6 +112,7 @@ def test_ngram_proposer():
112 # In this case, the proposer should return the 4-gram match.
113 token_ids_cpu = np.array([[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]])
114 result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose(
115 + num_speculative_tokens=2,
116 sampled_token_ids=[[0]],
117 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
118 token_ids_cpu=token_ids_cpu,
@@ -118,6 +122,7 @@ def test_ngram_proposer():
122 # Match for 2-gram and 3-gram, but not 4-gram.
123 token_ids_cpu = np.array([[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]])
124 result = get_ngram_proposer(min_n=2, max_n=4, k=2).propose(
125 + num_speculative_tokens=2,
126 sampled_token_ids=[[0]],
127 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
128 token_ids_cpu=token_ids_cpu,
@@ -127,6 +132,7 @@ def test_ngram_proposer():
132 # Multiple 3-gram matched, but always pick the first one.
133 token_ids_cpu = np.array([[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]])
134 result = get_ngram_proposer(min_n=3, max_n=3, k=2).propose(
135 + num_speculative_tokens=2,
136 sampled_token_ids=[[0]],
137 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
138 token_ids_cpu=token_ids_cpu,
@@ -136,6 +142,7 @@ def test_ngram_proposer():
142 # check empty input
143 token_ids_cpu = np.array([[]])
144 result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
145 + num_speculative_tokens=2,
146 sampled_token_ids=[[0]],
147 num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]),
148 token_ids_cpu=token_ids_cpu,
@@ -147,6 +154,7 @@ def test_ngram_proposer():
154 # second request has 3 tokens and no match. Padded with -1 for max len 5
155 token_ids_cpu = np.array([[1, 2, 3, 1, 2], [4, 5, 6, -1, -1]])
156 result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose(
157 + num_speculative_tokens=2,
158 sampled_token_ids=[[0], [1]],
159 num_tokens_no_spec=np.array([5, 3]),
160 token_ids_cpu=token_ids_cpu,
@@ -166,6 +174,7 @@ def test_ngram_proposer():
174 num_tokens_no_spec = np.array([5, 3, 5], dtype=np.int32)
175 sampled_token_ids = [[2], [], [8]] # Empty list for request 1 simulates prefill
176 result = proposer.propose(
177 + num_speculative_tokens=2,
178 sampled_token_ids=sampled_token_ids,
179 num_tokens_no_spec=num_tokens_no_spec,
180 token_ids_cpu=token_ids_cpu,
@@ -195,6 +204,7 @@ def test_ngram_proposer():
204 input_2[:3] = [4, 5, 6]
205 token_ids_cpu = np.array([input_1, input_2])
206 result = ngram_proposer.propose(
207 + num_speculative_tokens=2,
208 sampled_token_ids=[[0], [1]],
209 num_tokens_no_spec=np.array([len(input_1), 3]),
210 token_ids_cpu=token_ids_cpu,
vllm/config/speculative.py
+11
@@ -157,6 +157,14 @@ class SpeculativeConfig:
157 target_parallel_config: SkipValidation[ParallelConfig] = None # type: ignore
158 """The parallel configuration for the target model."""
159
160 + # dynamic speculative decoding control
161 + num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None
162 + """Batch-size schedule used to dynamically choose speculative-token count.
163 +
164 + Each entry is ``(range_start, range_end, num_speculative_tokens)`` with an
165 + inclusive batch-size range.
166 + """
167 +
168 # params generated in the post-init stage
169 draft_model_config: SkipValidation[ModelConfig] = None # type: ignore
170 """The configuration of the draft model initialized internal."""
@@ -1073,6 +1081,9 @@ class SpeculativeConfig:
1081 def use_dflash(self) -> bool:
1082 return self.method == "dflash"
1083
1084 + def uses_dynamic_speculative_decoding(self) -> bool:
1085 + return self.num_speculative_tokens_per_batch_size is not None
1086 +
1087 def uses_draft_model(self) -> bool:
1088 return self.method == "draft_model"
1089
vllm/config/vllm.py
+22
@@ -764,6 +764,23 @@ class VllmConfig:
764
765 apply_recursive(self, defaults)
766
767 + def _maybe_override_dynamic_sd_cudagraph_mode(self) -> None:
768 + speculative_config = self.speculative_config
769 + if (
770 + speculative_config is None
771 + or not speculative_config.uses_dynamic_speculative_decoding()
772 + or not self.compilation_config.cudagraph_mode.has_full_cudagraphs()
773 + ):
774 + return
775 +
776 + logger.warning_once(
777 + "Dynamic speculative decoding changes the target verification "
778 + "length at runtime. Overriding cudagraph_mode from %s to "
779 + "PIECEWISE for reliability.",
780 + self.compilation_config.cudagraph_mode.name,
781 + )
782 + self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
783 +
784 def _post_init_kv_transfer_config(self) -> None:
785 """Update KVTransferConfig based on top-level configs in VllmConfig.
786
@@ -1153,6 +1170,8 @@ class VllmConfig:
1170 "optimization level defaults."
1171 )
1172
1173 + self._maybe_override_dynamic_sd_cudagraph_mode()
1174 +
1175 if (
1176 self.compilation_config.cudagraph_mode.requires_piecewise_compilation()
1177 and self.compilation_config.mode != CompilationMode.VLLM_COMPILE
@@ -2005,6 +2024,9 @@ class VllmConfig:
2024 elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"):
2025 unsupported.append(f"speculative method '{speculative_config.method}'")
2026
2027 + if speculative_config.uses_dynamic_speculative_decoding():
2028 + unsupported.append("dynamic speculative decoding")
2029 +
2030 # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle)
2031 # DFlash uses parallel drafting natively in V2 via DFlashSpeculator.
2032 if (
vllm/v1/core/sched/async_scheduler.py
+4
@@ -19,6 +19,10 @@ class AsyncScheduler(Scheduler):
19 def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
20 super()._update_after_schedule(scheduler_output)
21 spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens
22 + # Use the latest num of scheduled draft tokens in next step as placeholder.
23 + self._spec_token_placeholders = [
24 + -1
25 + ] * scheduler_output.num_spec_tokens_to_schedule
26 for req_id in scheduler_output.num_scheduled_tokens:
27 request = self.requests[req_id]
28 if request.is_prefill_chunk:
vllm/v1/core/sched/output.py
+4
@@ -240,6 +240,10 @@ class SchedulerOutput:
240 # preventing stale NaN/data from corrupting attention or SSM computation.
241 new_block_ids_to_zero: list[int] | None = None
242
243 + # Dynamic speculative decoding: optimal K chosen by scheduler.
244 + # Number of spec tokens to schedule for the next step.
245 + num_spec_tokens_to_schedule: int = 0
246 +
247 @classmethod
248 def make_empty(cls) -> "SchedulerOutput":
249 return cls(
vllm/v1/core/sched/scheduler.py
+16
@@ -56,6 +56,7 @@ from vllm.v1.metrics.perf import ModelMetrics, PerfStats
56 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
57 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
58 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
59 +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
60 from vllm.v1.spec_decode.metrics import SpecDecodingStats
61 from vllm.v1.structured_output import StructuredOutputManager
62 from vllm.v1.utils import record_function_or_nullcontext
@@ -218,7 +219,14 @@ class Scheduler(SchedulerInterface):
219 self.use_eagle = False
220 self.num_spec_tokens = vllm_config.num_speculative_tokens
221 self.num_lookahead_tokens = 0
222 + self.dynamic_sd_lookup: list[int] | None = None
223 if speculative_config is not None:
224 + if speculative_config.num_speculative_tokens_per_batch_size:
225 + self.dynamic_sd_lookup = build_dynamic_sd_schedule_lookup(
226 + speculative_config.num_speculative_tokens_per_batch_size,
227 + vllm_max_batch_size=self.scheduler_config.max_num_seqs,
228 + vllm_num_speculative_tokens=self.num_spec_tokens,
229 + )
230 if speculative_config.use_eagle():
231 self.use_eagle = True
232 self.num_lookahead_tokens = self.num_spec_tokens
@@ -995,6 +1003,13 @@ class Scheduler(SchedulerInterface):
1003 else None
1004 )
1005
1006 + # Dynamic speculative decoding: compute optimal K
1007 + num_spec_tokens_to_schedule = self.num_spec_tokens
1008 + if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0:
1009 + num_spec_tokens_to_schedule = self.dynamic_sd_lookup[
1010 + len(num_scheduled_tokens)
1011 + ]
1012 +
1013 scheduler_output = SchedulerOutput(
1014 scheduled_new_reqs=new_reqs_data,
1015 scheduled_cached_reqs=cached_reqs_data,
@@ -1011,6 +1026,7 @@ class Scheduler(SchedulerInterface):
1026 finished_req_ids=self.finished_req_ids,
1027 free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
1028 new_block_ids_to_zero=new_block_ids_to_zero,
1029 + num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
1030 )
1031
1032 # NOTE(Kuntai): this function is designed for multiple purposes:
vllm/v1/spec_decode/dynamic/__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/v1/spec_decode/dynamic/utils.py new
+148
@@ -0,0 +1,148 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +DynamicSDSchedule = list[tuple[int, int, int]]
5 +
6 +
7 +def validate_and_normalize_dynamic_sd_schedule(
8 + num_speculative_tokens_per_batch_size: object,
9 +) -> DynamicSDSchedule:
10 + """Validate and normalize a Dynamic SD batch-size schedule.
11 +
12 + The schedule is expressed as a list of inclusive ranges:
13 +
14 + ``[(range_start, range_end, num_speculative_tokens), ...]``
15 + """
16 + if num_speculative_tokens_per_batch_size is None:
17 + raise ValueError(
18 + "num_speculative_tokens_per_batch_size is required for "
19 + "dynamic speculative decoding."
20 + )
21 + if not isinstance(num_speculative_tokens_per_batch_size, list):
22 + raise ValueError(
23 + "num_speculative_tokens_per_batch_size must be a non-empty list of "
24 + "(range_start, range_end, num_speculative_tokens) entries."
25 + )
26 + if not num_speculative_tokens_per_batch_size:
27 + raise ValueError("num_speculative_tokens_per_batch_size must not be empty.")
28 +
29 + parsed_schedule: DynamicSDSchedule = []
30 + for entry in num_speculative_tokens_per_batch_size:
31 + if not isinstance(entry, list | tuple) or len(entry) != 3:
32 + raise ValueError(
33 + "Each num_speculative_tokens_per_batch_size entry must be a "
34 + "3-item sequence: (range_start, range_end, num_speculative_tokens)."
35 + )
36 +
37 + range_start, range_end, num_speculative_tokens = (
38 + int(entry[0]),
39 + int(entry[1]),
40 + int(entry[2]),
41 + )
42 +
43 + if range_start <= 0 or range_end <= 0:
44 + raise ValueError(
45 + f"Batch-size range ({range_start}, {range_end}) must be positive."
46 + )
47 + if range_start > range_end:
48 + raise ValueError(
49 + "Batch-size range start must be <= end for "
50 + f"({range_start}, {range_end}, {num_speculative_tokens})."
51 + )
52 + if num_speculative_tokens < 0:
53 + raise ValueError(
54 + "num_speculative_tokens_per_batch_size values must be >= 0."
55 + )
56 +
57 + parsed_schedule.append((range_start, range_end, num_speculative_tokens))
58 +
59 + parsed_schedule.sort(key=lambda entry: entry[0])
60 +
61 + previous_end = 0
62 + for range_start, range_end, _ in parsed_schedule:
63 + if range_start <= previous_end:
64 + raise ValueError("Batch-size ranges must be non-overlapping and sorted.")
65 + previous_end = range_end
66 +
67 + first_range_start = parsed_schedule[0][0]
68 + if first_range_start != 1:
69 + raise ValueError(
70 + "The first batch-size range must start at 1 so every runtime "
71 + "batch size has a defined schedule."
72 + )
73 +
74 + return parsed_schedule
75 +
76 +
77 +def build_dynamic_sd_schedule_lookup(
78 + num_speculative_tokens_per_batch_size: object,
79 + vllm_max_batch_size: int,
80 + vllm_num_speculative_tokens: int,
81 +) -> list[int]:
82 + """Expand the configured schedule into a dense batch_size -> K lookup.
83 +
84 + "dense_schedule" means a 1-indexed lookup table where index ``batch_size``
85 + stores the exact K to use for that runtime batch size. This lets the
86 + scheduler do a simple array lookup instead of searching the configured
87 + ranges on every scheduling step.
88 + """
89 + if vllm_max_batch_size <= 0:
90 + raise ValueError("vllm_max_batch_size must be > 0.")
91 + if vllm_num_speculative_tokens <= 0:
92 + raise ValueError("vllm_num_speculative_tokens must be > 0.")
93 +
94 + parsed_schedule = validate_and_normalize_dynamic_sd_schedule(
95 + num_speculative_tokens_per_batch_size
96 + )
97 +
98 + # Index 0 is intentionally unused so that valid runtime batch sizes can be
99 + # looked up directly as dense_schedule[batch_size].
100 + dense_schedule = [0] * (vllm_max_batch_size + 1)
101 + next_batch_size = 1
102 + last_num_speculative_tokens: int | None = None
103 +
104 + for range_start, range_end, num_speculative_tokens in parsed_schedule:
105 + if range_start > next_batch_size and last_num_speculative_tokens is not None:
106 + # Fill any gap before the next configured range by carrying forward
107 + # the previous K. For example, [(1, 16, 3), (32, 128, 2)] should map
108 + # batch sizes 17-31 to K=3.
109 + for batch_size in range(
110 + next_batch_size,
111 + min(range_start, vllm_max_batch_size + 1),
112 + ):
113 + dense_schedule[batch_size] = min(
114 + vllm_num_speculative_tokens,
115 + last_num_speculative_tokens,
116 + )
117 +
118 + # Fill the current configured inclusive range with its K value.
119 + for batch_size in range(
120 + max(range_start, next_batch_size),
121 + min(range_end, vllm_max_batch_size) + 1,
122 + ):
123 + dense_schedule[batch_size] = min(
124 + vllm_num_speculative_tokens,
125 + num_speculative_tokens,
126 + )
127 +
128 + next_batch_size = max(next_batch_size, range_end + 1)
129 + last_num_speculative_tokens = num_speculative_tokens
130 +
131 + if next_batch_size > vllm_max_batch_size:
132 + break
133 +
134 + if last_num_speculative_tokens is None:
135 + raise ValueError(
136 + "num_speculative_tokens_per_batch_size must contain at least "
137 + "one valid batch-size range."
138 + )
139 +
140 + # Fill the tail after the final configured range by carrying forward the
141 + # last K through vllm_max_batch_size.
142 + for batch_size in range(next_batch_size, vllm_max_batch_size + 1):
143 + dense_schedule[batch_size] = min(
144 + vllm_num_speculative_tokens,
145 + last_num_speculative_tokens,
146 + )
147 +
148 + return dense_schedule
vllm/v1/spec_decode/extract_hidden_states.py
+6 -1
@@ -29,7 +29,10 @@ class ExtractHiddenStatesProposer:
29 def __init__(self, vllm_config: VllmConfig, device):
30 assert vllm_config.speculative_config is not None
31
32 - assert vllm_config.speculative_config.num_speculative_tokens == 1
32 + self.num_speculative_tokens = (
33 + vllm_config.speculative_config.num_speculative_tokens
34 + )
35 + assert self.num_speculative_tokens == 1
36 if vllm_config.speculative_config.disable_padded_drafter_batch:
37 raise ValueError(
38 "disable_padded_drafter_batch is not supported with "
@@ -82,6 +85,7 @@ class ExtractHiddenStatesProposer:
85
86 def propose(
87 self,
88 + num_speculative_tokens: int,
89 sampled_token_ids: torch.Tensor,
90 target_hidden_states: list[torch.Tensor],
91 common_attn_metadata: CommonAttentionMetadata,
@@ -112,6 +116,7 @@ class ExtractHiddenStatesProposer:
116 - Draft tokens matching sampled tokens, shape [batch_size, 1]
117 - KV connector output (if KV transfer is active), else None
118 """
119 + assert num_speculative_tokens == self.num_speculative_tokens
120 assert self.model is not None and isinstance(target_hidden_states, list)
121
122 # target_hidden_states is a list of tensors (one per layer)
vllm/v1/spec_decode/llm_base_proposer.py
+13
@@ -434,6 +434,7 @@ class SpecDecodeBaseProposer:
434
435 def propose(
436 self,
437 + num_speculative_tokens,
438 # [num_tokens]
439 target_token_ids: torch.Tensor,
440 # [num_tokens] or [3, num_tokens] when M-RoPE is enabled
@@ -451,6 +452,7 @@ class SpecDecodeBaseProposer:
452 | list[dict[str, torch.Tensor]]
453 | None = None,
454 ) -> torch.Tensor:
455 + self.num_speculative_tokens = num_speculative_tokens
456 self._last_draft_probs = None
457 batch_size = common_attn_metadata.batch_size()
458
@@ -521,6 +523,17 @@ class SpecDecodeBaseProposer:
523
524 sample_hidden_states = last_hidden_states[token_indices_to_sample]
525
526 + # No draft tokens requested (e.g. Dynamic SD decided K=0).
527 + # The prefill forward pass above already ran to keep the drafter
528 + # KV cache in sync, so just return an empty tensor.
529 + if self.num_speculative_tokens == 0:
530 + return torch.empty(
531 + batch_size,
532 + 0,
533 + device=sample_hidden_states.device,
534 + dtype=torch.int64,
535 + )
536 +
537 # Early exit if there is only one draft token to be generated.
538 if self.num_speculative_tokens == 1 or self.parallel_drafting:
539 draft_token_ids, draft_probs = self._sample_draft_tokens(
vllm/v1/spec_decode/medusa.py
+3
@@ -35,15 +35,18 @@ class MedusaProposer:
35 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
36 self.hidden_size = self.spec_config.draft_model_config.get_hidden_size()
37 self.dtype = vllm_config.model_config.dtype
38 + self.num_speculative_tokens = self.spec_config.num_speculative_tokens
39
40 def propose(
41 self,
42 + num_speculative_tokens: int,
43 target_hidden_states: torch.Tensor,
44 sampling_metadata: SamplingMetadata,
45 slot_mappings: dict[str, torch.Tensor]
46 | list[dict[str, torch.Tensor]]
47 | None = None, # unused
48 ) -> torch.Tensor:
49 + assert num_speculative_tokens == self.num_speculative_tokens
50 # Generate blocks and compute logits
51 blocks = self.model(target_hidden_states)
52 logits = self.model.compute_logits(blocks)
vllm/v1/spec_decode/metrics.py
+4
@@ -28,12 +28,14 @@ class SpecDecodingStats:
28 num_draft_tokens: int = 0
29 num_accepted_tokens: int = 0
30 num_accepted_tokens_per_pos: list[int] = field(default_factory=list)
31 + num_draft_tokens_per_pos: list[int] = field(default_factory=list)
32
33 @classmethod
34 def new(cls, num_spec_tokens: int) -> "SpecDecodingStats":
35 return cls(
36 num_spec_tokens=num_spec_tokens,
37 num_accepted_tokens_per_pos=[0] * num_spec_tokens,
38 + num_draft_tokens_per_pos=[0] * num_spec_tokens,
39 )
40
41 def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int):
@@ -43,6 +45,8 @@ class SpecDecodingStats:
45 assert num_accepted_tokens <= self.num_spec_tokens
46 for i in range(num_accepted_tokens):
47 self.num_accepted_tokens_per_pos[i] += 1
48 + for i in range(num_draft_tokens):
49 + self.num_draft_tokens_per_pos[i] += 1
50
51
52 class SpecDecodingLogging:
vllm/v1/spec_decode/ngram_proposer.py
+9 -1
@@ -55,6 +55,7 @@ class NgramProposer:
55 # Trigger Numba JIT compilation for N-gram proposer.
56 # This usually takes less than 1 second.
57 self.propose(
58 + self.k,
59 [[]] * 1024,
60 np.zeros(1024, dtype=np.int32),
61 np.zeros((1024, self.max_model_len), dtype=np.int32),
@@ -66,6 +67,7 @@ class NgramProposer:
67 valid_ngram_requests: list,
68 num_tokens_no_spec: np.ndarray,
69 token_ids_cpu: np.ndarray,
70 + k: int,
71 ) -> list[list[int]]:
72 """Batch version of ngram proposer using numba for acceleration.
73
@@ -78,6 +80,8 @@ class NgramProposer:
80 token_ids_cpu:
81 Numpy array of shape (batch_size, max_model_len)
82 representing the token IDs for each request.
83 + k:
84 + Number of speculative tokens to propose.
85
86 Returns:
87 list[list[int]]:
@@ -110,7 +114,7 @@ class NgramProposer:
114 self.min_n,
115 self.max_n,
116 self.max_model_len,
113 - self.k,
117 + k,
118 self.valid_ngram_draft,
119 self.valid_ngram_num_drafts,
120 )
@@ -130,6 +134,7 @@ class NgramProposer:
134
135 def propose(
136 self,
137 + num_speculative_tokens: int,
138 sampled_token_ids: list[list[int]],
139 num_tokens_no_spec: np.ndarray,
140 token_ids_cpu: np.ndarray,
@@ -137,6 +142,8 @@ class NgramProposer:
142 | list[dict[str, torch.Tensor]]
143 | None = None, # unused
144 ) -> list[list[int]]:
145 + assert num_speculative_tokens <= self.k
146 +
147 # find which requests need ngram proposals
148 valid_ngram_requests = []
149 for i, sampled_ids in enumerate(sampled_token_ids):
@@ -157,6 +164,7 @@ class NgramProposer:
164 valid_ngram_requests,
165 num_tokens_no_spec,
166 token_ids_cpu,
167 + num_speculative_tokens,
168 )
169
170 return draft_token_ids
vllm/v1/spec_decode/ngram_proposer_gpu.py
+3
@@ -314,6 +314,7 @@ class NgramProposerGPU:
314
315 def propose(
316 self,
317 + num_speculative_tokens: int,
318 num_tokens_no_spec: torch.Tensor, # [batch_size]
319 token_ids_gpu: torch.Tensor, # [batch_size, max_len]
320 valid_sampled_token_ids_gpu: torch.Tensor, # [batch_size, num_spec_tokens + 1]
@@ -326,6 +327,7 @@ class NgramProposerGPU:
327 updated lengths, then run the kernel.
328
329 Args:
330 + num_speculative_tokens: Number of speculative tokens to propose.
331 num_tokens_no_spec: Number of tokens per sequence (read-only)
332 token_ids_gpu: Token IDs tensor (modified in-place with new tokens)
333 valid_sampled_token_ids_gpu: Newly sampled tokens to scatter
@@ -336,6 +338,7 @@ class NgramProposerGPU:
338 num_valid_draft_tokens: Count of leading valid draft tokens
339 per request [batch_size]
340 """
341 + assert num_speculative_tokens == self.k
342 assert token_ids_gpu.device == self.device
343 assert num_tokens_no_spec.device == self.device
344
vllm/v1/spec_decode/step3p5.py
+2
@@ -273,6 +273,7 @@ class Step3p5MTPProposer(EagleProposer):
273
274 def propose(
275 self,
276 + num_speculative_tokens: int,
277 target_token_ids: torch.Tensor,
278 target_positions: torch.Tensor,
279 target_hidden_states: torch.Tensor,
@@ -286,6 +287,7 @@ class Step3p5MTPProposer(EagleProposer):
287 | list[dict[str, torch.Tensor]]
288 | None = None,
289 ) -> torch.Tensor:
290 + self.num_speculative_tokens = num_speculative_tokens
291 self._last_draft_probs = None
292 batch_size = common_attn_metadata.batch_size()
293
vllm/v1/spec_decode/suffix_decoding.py
+2
@@ -34,12 +34,14 @@ class SuffixDecodingProposer:
34
35 def propose(
36 self,
37 + num_speculative_tokens: int,
38 input_batch: InputBatch,
39 sampled_token_ids: list[list[int]],
40 slot_mappings: dict[str, torch.Tensor]
41 | list[dict[str, torch.Tensor]]
42 | None = None, # unused
43 ) -> list[list[int]]:
44 + assert num_speculative_tokens == self.num_speculative_tokens
45 """
46 Propose speculative tokens for each request in the input batch. Suffix Decoding
47 will speculate a dynamic number of tokens for each request every decoding step,
vllm/v1/worker/gpu_model_runner.py
+24 -5
@@ -620,9 +620,11 @@ class GPUModelRunner(
620 )
621
622 self.num_spec_tokens = 0
623 + self.prev_num_spec_tokens = 0
624 self.valid_sampled_token_count_gpu: torch.Tensor | None = None
625 if self.speculative_config:
626 self.num_spec_tokens = self.speculative_config.num_speculative_tokens
627 + self.prev_num_spec_tokens = self.num_spec_tokens
628 draft_config = self.speculative_config.draft_model_config
629 if draft_config is not None and draft_config.max_model_len is not None:
630 self.effective_drafter_max_model_len = draft_config.max_model_len
@@ -1764,7 +1766,7 @@ class GPUModelRunner(
1766 spec_flattened_indices.extend(
1767 range(flattened_index - draft_len + 1, flattened_index + 1)
1768 )
1767 - start = prev_index * self.num_spec_tokens
1769 + start = prev_index * self.prev_num_spec_tokens
1770 # prev_draft_token_indices is used to find which draft_tokens_id
1771 # should be copied to input_ids
1772 # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]]
@@ -4704,6 +4706,9 @@ class GPUModelRunner(
4706 def _copy_draft_token_ids_to_cpu(
4707 self, scheduler_output: "SchedulerOutput", zeros_only: bool = False
4708 ) -> None:
4709 + if torch.is_tensor(self._draft_token_ids):
4710 + assert isinstance(self._draft_token_ids, torch.Tensor)
4711 + self.prev_num_spec_tokens = self._draft_token_ids.shape[1]
4712 # Check if we need to copy draft tokens to CPU. In async scheduling,
4713 # we only copy when needed for structured output, penalties or bad_words.
4714 if self.use_async_scheduling and not (
@@ -4722,16 +4727,17 @@ class GPUModelRunner(
4727 assert self.draft_token_ids_cpu is not None
4728 default_stream = torch.cuda.current_stream()
4729 num_reqs = draft_token_ids.shape[0]
4730 + num_spec_tokens = draft_token_ids.shape[1]
4731 with torch.cuda.stream(self.draft_token_ids_copy_stream):
4732 if not zeros_only:
4733 # Trigger async copy of draft token ids to cpu.
4734 self.draft_token_ids_copy_stream.wait_stream(default_stream)
4729 - self.draft_token_ids_cpu[:num_reqs].copy_(
4735 + self.draft_token_ids_cpu[:num_reqs, :num_spec_tokens].copy_(
4736 draft_token_ids, non_blocking=True
4737 )
4738 else:
4739 # No copy needed, just zero-out cpu tensor.
4734 - self.draft_token_ids_cpu[:num_reqs] = 0
4740 + self.draft_token_ids_cpu[:num_reqs, :num_spec_tokens] = 0
4741 self.draft_token_ids_event.record()
4742
4743 def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]:
@@ -4743,7 +4749,11 @@ class GPUModelRunner(
4749 assert self.draft_token_ids_event is not None
4750 assert self.draft_token_ids_cpu is not None
4751 self.draft_token_ids_event.synchronize()
4746 - return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids
4752 + assert isinstance(self._draft_token_ids, torch.Tensor)
4753 + num_spec_tokens = self._draft_token_ids.shape[1]
4754 + return self.draft_token_ids_cpu[
4755 + : len(req_ids), :num_spec_tokens
4756 + ].tolist(), req_ids
4757
4758 def _copy_valid_sampled_token_count(
4759 self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor
@@ -4823,6 +4833,7 @@ class GPUModelRunner(
4833 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
4834 spec_config = self.speculative_config
4835 assert spec_config is not None
4836 + num_spec_tokens_to_schedule = scheduler_output.num_spec_tokens_to_schedule
4837 self._draft_probs = None
4838 self._draft_prob_req_ids = None
4839 if spec_config.method == "ngram":
@@ -4831,6 +4842,7 @@ class GPUModelRunner(
4842 assert isinstance(sampled_token_ids, list)
4843 assert isinstance(self.drafter, NgramProposer)
4844 draft_token_ids = self.drafter.propose(
4845 + num_spec_tokens_to_schedule,
4846 sampled_token_ids,
4847 self.input_batch.num_tokens_no_spec,
4848 self.input_batch.token_ids_cpu,
@@ -4864,6 +4876,7 @@ class GPUModelRunner(
4876 batch_size = next_token_ids.shape[0]
4877
4878 draft_token_ids, num_valid_draft_tokens = self.drafter.propose(
4879 + num_spec_tokens_to_schedule,
4880 self.num_tokens_no_spec_gpu[:batch_size],
4881 self.token_ids_gpu_tensor[:batch_size],
4882 valid_sampled_token_ids_gpu,
@@ -4885,7 +4898,10 @@ class GPUModelRunner(
4898 assert isinstance(sampled_token_ids, list)
4899 assert isinstance(self.drafter, SuffixDecodingProposer)
4900 draft_token_ids = self.drafter.propose(
4888 - self.input_batch, sampled_token_ids, slot_mappings=slot_mappings
4901 + num_spec_tokens_to_schedule,
4902 + self.input_batch,
4903 + sampled_token_ids,
4904 + slot_mappings=slot_mappings,
4905 )
4906 elif spec_config.method == "medusa":
4907 assert isinstance(sampled_token_ids, list)
@@ -4909,6 +4925,7 @@ class GPUModelRunner(
4925 hidden_states = sample_hidden_states[indices]
4926
4927 draft_token_ids = self.drafter.propose(
4928 + num_speculative_tokens=num_spec_tokens_to_schedule,
4929 target_hidden_states=hidden_states,
4930 sampling_metadata=sampling_metadata,
4931 slot_mappings=slot_mappings,
@@ -4926,6 +4943,7 @@ class GPUModelRunner(
4943 target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states]
4944
4945 draft_token_ids = self.drafter.propose(
4946 + num_speculative_tokens=num_spec_tokens_to_schedule,
4947 sampled_token_ids=sampled_token_ids,
4948 target_hidden_states=target_hidden_states,
4949 common_attn_metadata=common_attn_metadata,
@@ -5059,6 +5077,7 @@ class GPUModelRunner(
5077 mm_embed_inputs = None
5078
5079 draft_token_ids = self.drafter.propose(
5080 + num_speculative_tokens=num_spec_tokens_to_schedule,
5081 target_token_ids=target_token_ids,
5082 target_positions=target_positions,
5083 target_hidden_states=target_hidden_states,