@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 72c0d67657

[Feat] Support thinking_token_budget in Model Runner V2 (#46727)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Chauncey committed Aug 7, 2026 at 11:36 UTC 72c0d6765793e4c7242c3586274af3e1a8aca170
10 files changed +1122 -43
benchmarks/kernels/benchmark_thinking_budget.py new
+365
@@ -0,0 +1,365 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import dataclasses
5 +from collections.abc import Callable
6 +from typing import Literal
7 +
8 +import numpy as np
9 +import torch
10 +
11 +from vllm.sampling_params import SamplingParams
12 +from vllm.triton_utils import triton
13 +from vllm.utils.argparse_utils import FlexibleArgumentParser
14 +from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
15 +from vllm.v1.worker.gpu.states import RequestState
16 +
17 +START_TOKEN_ID = 90
18 +END_TOKEN_ID = 91
19 +TOKEN_ID = 10
20 +VOCAB_SIZE = 128
21 +QUANTILES = [0.5, 0.2, 0.8]
22 +
23 +
24 +class ReasoningConfig:
25 + reasoning_start_token_ids = [START_TOKEN_ID]
26 + reasoning_end_token_ids = [END_TOKEN_ID]
27 + natural_reasoning_end_token_ids = [END_TOKEN_ID]
28 +
29 +
30 +@dataclasses.dataclass
31 +class BenchmarkCase:
32 + req_states: RequestState
33 + state: ThinkingBudgetState
34 + run: Callable[[], None]
35 +
36 +
37 +def create_case(
38 + history_len: int,
39 + device: torch.device,
40 + extra_tokens: int = 0,
41 + history_type: Literal["reasoning", "prefill"] = "reasoning",
42 + batch_size: int = 1,
43 + budget_type: Literal["active", "forced", "mixed"] = "active",
44 +) -> BenchmarkCase:
45 + assert history_len >= 2
46 + if history_type == "prefill":
47 + tokens = [*([TOKEN_ID] * (history_len - 1)), START_TOKEN_ID]
48 + else:
49 + tokens = [TOKEN_ID, START_TOKEN_ID, *([TOKEN_ID] * (history_len - 2))]
50 + req_states = RequestState(
51 + max_num_reqs=batch_size,
52 + max_model_len=history_len + extra_tokens + 1,
53 + max_num_batched_tokens=batch_size,
54 + num_speculative_steps=1,
55 + vocab_size=VOCAB_SIZE,
56 + device=device,
57 + )
58 + req_indices = []
59 + for req_num in range(batch_size):
60 + req_id = f"benchmark-{req_num}"
61 + req_states.add_request(
62 + req_id=req_id,
63 + prompt_len=1,
64 + all_token_ids=tokens,
65 + num_computed_tokens=history_len,
66 + max_tokens=extra_tokens + 1,
67 + )
68 + req_indices.append(req_states.req_id_to_index[req_id])
69 + req_states.apply_staged_writes()
70 +
71 + state = ThinkingBudgetState(req_states, ReasoningConfig())
72 + active_budget = history_len + extra_tokens + 1
73 + forced_budget = max(1, history_len - 2)
74 + for req_num, req_idx in enumerate(req_indices):
75 + if budget_type == "forced":
76 + budget = forced_budget
77 + elif budget_type == "mixed" and req_num % 2 == 1:
78 + budget = None
79 + else:
80 + budget = active_budget
81 + state.add_request(
82 + req_idx,
83 + SamplingParams(thinking_token_budget=budget),
84 + )
85 + state.apply_staged_writes()
86 +
87 + idx_mapping = torch.tensor(req_indices, dtype=torch.int32, device=device)
88 + logits = torch.zeros((batch_size, VOCAB_SIZE), device=device)
89 + idx_mapping_np = idx_mapping.cpu().numpy()
90 + input_ids = torch.full((batch_size,), TOKEN_ID, dtype=torch.int32, device=device)
91 + local_pos = torch.zeros(batch_size, dtype=torch.int32, device=device)
92 +
93 + def run() -> None:
94 + state.apply(
95 + logits,
96 + idx_mapping,
97 + idx_mapping,
98 + idx_mapping_np,
99 + input_ids,
100 + local_pos,
101 + )
102 +
103 + run()
104 + torch.accelerator.synchronize(device)
105 + active_req_indices = [
106 + req_idx
107 + for req_num, req_idx in enumerate(req_indices)
108 + if budget_type != "mixed" or req_num % 2 == 0
109 + ]
110 + assert torch.all(state.cached_scan_pos[active_req_indices] == history_len).item()
111 + if budget_type == "forced":
112 + assert torch.all(logits[:, END_TOKEN_ID] == 1.0e9).item()
113 + return BenchmarkCase(req_states, state, run)
114 +
115 +
116 +def benchmark_cached(
117 + case: BenchmarkCase, warmup_ms: int, rep_ms: int
118 +) -> tuple[float, float, float]:
119 + median_ms, min_ms, max_ms = triton.testing.do_bench(
120 + case.run,
121 + warmup=warmup_ms,
122 + rep=rep_ms,
123 + quantiles=QUANTILES,
124 + )
125 + return median_ms, min_ms, max_ms
126 +
127 +
128 +def summarize(timings_ms: list[float]) -> tuple[float, float, float]:
129 + median_ms, min_ms, max_ms = np.quantile(timings_ms, QUANTILES)
130 + return float(median_ms), float(min_ms), float(max_ms)
131 +
132 +
133 +def time_cuda_call(run: Callable[[], None]) -> float:
134 + start = torch.Event(enable_timing=True)
135 + end = torch.Event(enable_timing=True)
136 + start.record()
137 + run()
138 + end.record()
139 + end.synchronize()
140 + return start.elapsed_time(end)
141 +
142 +
143 +def benchmark_incremental_decode(
144 + history_len: int,
145 + device: torch.device,
146 + warmup_steps: int,
147 + iterations: int,
148 +) -> tuple[float, float, float]:
149 + case = create_case(history_len, device, warmup_steps + iterations)
150 + timings_ms: list[float] = []
151 + total_len = history_len
152 + for step in range(warmup_steps + iterations):
153 + case.req_states.all_token_ids.stage_write(0, total_len, [TOKEN_ID])
154 + total_len += 1
155 + case.req_states.total_len.stage_write_elem(0, total_len)
156 + case.req_states.apply_staged_writes()
157 + torch.accelerator.synchronize(device)
158 +
159 + elapsed_ms = time_cuda_call(case.run)
160 + if step >= warmup_steps:
161 + timings_ms.append(elapsed_ms)
162 +
163 + assert case.state.cached_scan_pos[0].item() == total_len
164 + return summarize(timings_ms)
165 +
166 +
167 +def benchmark_cold_scan(
168 + history_len: int,
169 + device: torch.device,
170 + warmup_steps: int,
171 + iterations: int,
172 + history_type: Literal["reasoning", "prefill"],
173 +) -> tuple[float, float, float]:
174 + case = create_case(history_len, device, history_type=history_type)
175 + timings_ms: list[float] = []
176 + for step in range(warmup_steps + iterations):
177 + case.state.cached_last_start.fill_(-1)
178 + case.state.cached_last_end.fill_(-1)
179 + case.state.cached_scan_pos.zero_()
180 + torch.accelerator.synchronize(device)
181 +
182 + elapsed_ms = time_cuda_call(case.run)
183 + if step >= warmup_steps:
184 + timings_ms.append(elapsed_ms)
185 +
186 + return summarize(timings_ms)
187 +
188 +
189 +def main() -> None:
190 + parser = FlexibleArgumentParser(
191 + description=(
192 + "Benchmark thinking-budget scan, forced-end, and batched request "
193 + "overhead across token history lengths, including worst-case cold "
194 + "resume reconstruction."
195 + )
196 + )
197 + parser.add_argument(
198 + "--history-lengths",
199 + type=int,
200 + nargs="+",
201 + default=[128, 16384, 32768],
202 + )
203 + parser.add_argument(
204 + "--modes",
205 + nargs="+",
206 + choices=[
207 + "cached",
208 + "incremental-decode",
209 + "cold-prefill",
210 + "cold-resume-worst-case",
211 + "forced-end",
212 + "batched-budgeted",
213 + "batched-mixed",
214 + ],
215 + default=[
216 + "cached",
217 + "incremental-decode",
218 + "cold-prefill",
219 + "cold-resume-worst-case",
220 + "forced-end",
221 + "batched-budgeted",
222 + "batched-mixed",
223 + ],
224 + )
225 + parser.add_argument(
226 + "--batch-sizes",
227 + type=int,
228 + nargs="+",
229 + default=[1, 8, 32],
230 + help="Batch sizes used by batched-budgeted and batched-mixed modes.",
231 + )
232 + parser.add_argument("--device", type=str, default="cuda:0")
233 + parser.add_argument(
234 + "--warmup-ms",
235 + type=int,
236 + default=100,
237 + help="Warmup duration passed to triton.testing.do_bench.",
238 + )
239 + parser.add_argument(
240 + "--rep-ms",
241 + type=int,
242 + default=500,
243 + help="Measurement duration passed to triton.testing.do_bench.",
244 + )
245 + parser.add_argument(
246 + "--warmup-steps",
247 + type=int,
248 + default=10,
249 + help="Warmup steps for incremental-decode and cold modes.",
250 + )
251 + parser.add_argument(
252 + "--iterations",
253 + type=int,
254 + default=100,
255 + help="Measured steps for incremental-decode and cold modes.",
256 + )
257 + parser.add_argument(
258 + "--max-slowdown",
259 + type=float,
260 + default=None,
261 + help=(
262 + "Fail if cached or incremental-decode median latency exceeds this "
263 + "multiple of its shortest-history latency. Cold modes are excluded "
264 + "because worst-case cold-resume cost scales with scanned history."
265 + ),
266 + )
267 + args = parser.parse_args()
268 +
269 + history_lengths = sorted(set(args.history_lengths))
270 + if not history_lengths or history_lengths[0] < 2:
271 + parser.error("history lengths must be at least 2")
272 + if args.warmup_steps < 0 or args.iterations <= 0:
273 + parser.error("warmup steps must be non-negative and iterations positive")
274 + batch_sizes = sorted(set(args.batch_sizes))
275 + if not batch_sizes or batch_sizes[0] <= 0:
276 + parser.error("batch sizes must be positive")
277 +
278 + device = torch.device(args.device)
279 + results: dict[str, dict[int, tuple[float, float, float]]] = {}
280 + for mode in args.modes:
281 + mode_batch_sizes = batch_sizes if mode.startswith("batched-") else [1]
282 + if mode == "batched-mixed":
283 + mode_batch_sizes = [size for size in mode_batch_sizes if size >= 2]
284 + if not mode_batch_sizes:
285 + parser.error("batched-mixed requires a batch size of at least 2")
286 + for batch_size in mode_batch_sizes:
287 + label = f"{mode}-b{batch_size}" if mode.startswith("batched-") else mode
288 + mode_results: dict[int, tuple[float, float, float]] = {}
289 + for history_len in history_lengths:
290 + if mode == "cached":
291 + case = create_case(history_len, device)
292 + result = benchmark_cached(case, args.warmup_ms, args.rep_ms)
293 + elif mode == "incremental-decode":
294 + result = benchmark_incremental_decode(
295 + history_len,
296 + device,
297 + args.warmup_steps,
298 + args.iterations,
299 + )
300 + elif mode == "cold-prefill":
301 + result = benchmark_cold_scan(
302 + history_len,
303 + device,
304 + args.warmup_steps,
305 + args.iterations,
306 + "prefill",
307 + )
308 + elif mode == "cold-resume-worst-case":
309 + result = benchmark_cold_scan(
310 + history_len,
311 + device,
312 + args.warmup_steps,
313 + args.iterations,
314 + "reasoning",
315 + )
316 + else:
317 + budget_type = (
318 + "forced"
319 + if mode == "forced-end"
320 + else "mixed"
321 + if mode == "batched-mixed"
322 + else "active"
323 + )
324 + case = create_case(
325 + history_len,
326 + device,
327 + batch_size=batch_size,
328 + budget_type=budget_type,
329 + )
330 + result = benchmark_cached(case, args.warmup_ms, args.rep_ms)
331 + mode_results[history_len] = result
332 + results[label] = mode_results
333 +
334 + print("mode history_len median_us p20_us p80_us slowdown")
335 + for label, mode_results in results.items():
336 + baseline_ms = mode_results[history_lengths[0]][0]
337 + for history_len in history_lengths:
338 + median_ms, min_ms, max_ms = mode_results[history_len]
339 + slowdown = median_ms / baseline_ms
340 + print(
341 + f"{label:<26} {history_len:>11} {median_ms * 1000:>9.3f} "
342 + f"{min_ms * 1000:>6.3f} {max_ms * 1000:>6.3f} "
343 + f"{slowdown:>7.3f}x"
344 + )
345 +
346 + if args.max_slowdown is not None:
347 + for mode in ("cached", "incremental-decode"):
348 + if mode not in results:
349 + continue
350 + mode_results = results[mode]
351 + baseline_ms = mode_results[history_lengths[0]][0]
352 + worst_history_len = max(
353 + mode_results, key=lambda length: mode_results[length][0]
354 + )
355 + worst_slowdown = mode_results[worst_history_len][0] / baseline_ms
356 + if worst_slowdown > args.max_slowdown:
357 + raise SystemExit(
358 + f"thinking-budget {mode} slowdown {worst_slowdown:.3f}x "
359 + f"at history length {worst_history_len} exceeds "
360 + f"{args.max_slowdown:.3f}x"
361 + )
362 +
363 +
364 +if __name__ == "__main__":
365 + main()
tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
+3 -8
@@ -72,9 +72,7 @@ def server():
72 "0.4",
73 "--no-async-scheduling",
74 ]
75 - # thinking_token_budget is not yet supported by the V2 model runner.
76 - env_dict = {"VLLM_USE_V2_MODEL_RUNNER": "0"}
77 - with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
75 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
76 yield remote_server
77
78
@@ -90,9 +88,7 @@ def server_with_auto_reasoning_config():
88 "0.4",
89 "--no-async-scheduling",
90 ]
93 - # thinking_token_budget is not yet supported by the V2 model runner.
94 - env_dict = {"VLLM_USE_V2_MODEL_RUNNER": "0"}
95 - with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
91 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
92 yield remote_server
93
94
@@ -126,8 +122,7 @@ def server_qwen35_fp8_mtp_tp2():
122 }
123 ),
124 ]
129 - # thinking_token_budget is not yet supported by the V2 model runner.
130 - env_dict: dict[str, str] = {"VLLM_USE_V2_MODEL_RUNNER": "0"}
125 + env_dict: dict[str, str] = {}
126 # With 4+ GPUs, run TP=2 on physical devices 2,3 so module-scoped 0.6B servers
127 # on 0,1 do not exhaust memory on the same devices as this worker.
128 if current_platform.device_count() >= 4:
tests/v1/worker/test_gpu_thinking_budget.py new
+262
@@ -0,0 +1,262 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +pytest.importorskip("triton")
8 +if not torch.cuda.is_available():
9 + pytest.skip(
10 + "CUDA required for Model Runner V2 thinking budget tests",
11 + allow_module_level=True,
12 + )
13 +
14 +from vllm.sampling_params import SamplingParams
15 +from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
16 +from vllm.v1.worker.gpu.states import RequestState
17 +
18 +DEVICE = torch.device("cuda")
19 +START = 90
20 +END = 91
21 +END_A = 92
22 +END_B = 93
23 +VOCAB_SIZE = 128
24 +
25 +
26 +class MockReasoningConfig:
27 + reasoning_start_token_ids = [START]
28 + reasoning_end_token_ids = [END]
29 + natural_reasoning_end_token_ids = [END]
30 +
31 +
32 +class MockMultiTokenEndReasoningConfig:
33 + reasoning_start_token_ids = [START]
34 + reasoning_end_token_ids = [END_A, END_B]
35 + natural_reasoning_end_token_ids = [END_A, END_B]
36 +
37 +
38 +class MockDistinctEndReasoningConfig:
39 + reasoning_start_token_ids = [START]
40 + reasoning_end_token_ids = [END_A, END_B]
41 + natural_reasoning_end_token_ids = [END]
42 +
43 +
44 +def _make_req_states(tokens: list[int], prompt_len: int = 1) -> RequestState:
45 + req_states = RequestState(
46 + max_num_reqs=4,
47 + max_model_len=max(64, len(tokens) + 1),
48 + max_num_batched_tokens=16,
49 + num_speculative_steps=4,
50 + vocab_size=VOCAB_SIZE,
51 + device=DEVICE,
52 + )
53 + req_states.add_request(
54 + req_id="req",
55 + prompt_len=prompt_len,
56 + all_token_ids=tokens,
57 + num_computed_tokens=len(tokens),
58 + max_tokens=32,
59 + )
60 + req_states.apply_staged_writes()
61 + return req_states
62 +
63 +
64 +def _apply(
65 + state: ThinkingBudgetState,
66 + logits: torch.Tensor,
67 + input_ids: list[int],
68 + local_pos: list[int],
69 +) -> torch.Tensor:
70 + idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE)
71 + expanded_idx_mapping = torch.tensor(
72 + [3] * len(input_ids), dtype=torch.int32, device=DEVICE
73 + )
74 + idx_mapping_np = idx_mapping.cpu().numpy()
75 + state.apply(
76 + logits,
77 + expanded_idx_mapping,
78 + idx_mapping,
79 + idx_mapping_np,
80 + torch.tensor(input_ids, dtype=torch.int32, device=DEVICE),
81 + torch.tensor(local_pos, dtype=torch.int32, device=DEVICE),
82 + )
83 + return logits.cpu()
84 +
85 +
86 +def test_v2_thinking_budget_forces_end_after_budget_reached():
87 + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
88 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
89 + state.add_request(3, SamplingParams(thinking_token_budget=3))
90 + state.apply_staged_writes()
91 +
92 + logits = torch.arange(VOCAB_SIZE, dtype=torch.float32, device=DEVICE).view(1, -1)
93 + expected = logits.cpu()
94 + out = _apply(state, logits, input_ids=[12], local_pos=[0])
95 +
96 + expected[0, END] = 1.0e9
97 + torch.testing.assert_close(out, expected)
98 +
99 +
100 +def test_v2_thinking_budget_restores_masked_end_token():
101 + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
102 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
103 + state.add_request(3, SamplingParams(thinking_token_budget=3))
104 + state.apply_staged_writes()
105 +
106 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
107 + logits[0, END] = -float("inf")
108 + out = _apply(state, logits, input_ids=[12], local_pos=[0])
109 +
110 + assert out[0, END] == pytest.approx(1.0e9)
111 +
112 +
113 +def test_v2_thinking_budget_allows_tokens_before_budget():
114 + req_states = _make_req_states([1, START, 10, 11], prompt_len=1)
115 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
116 + state.add_request(3, SamplingParams(thinking_token_budget=3))
117 + state.apply_staged_writes()
118 +
119 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
120 + out = _apply(state, logits, input_ids=[11], local_pos=[0])
121 +
122 + assert torch.all(out == 0)
123 +
124 +
125 +def test_v2_thinking_budget_continues_multi_token_end_marker():
126 + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
127 + state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig())
128 + state.add_request(3, SamplingParams(thinking_token_budget=3))
129 + state.apply_staged_writes()
130 +
131 + logits = torch.zeros((2, VOCAB_SIZE), device=DEVICE)
132 + out = _apply(
133 + state,
134 + logits,
135 + input_ids=[12, END_A],
136 + local_pos=[0, 1],
137 + )
138 +
139 + assert out[0, END_A] == pytest.approx(1.0e9)
140 + assert out[1, END_B] == pytest.approx(1.0e9)
141 +
142 +
143 +def test_v2_thinking_budget_uses_distinct_forced_end_marker():
144 + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
145 + state = ThinkingBudgetState(req_states, MockDistinctEndReasoningConfig())
146 + state.add_request(3, SamplingParams(thinking_token_budget=3))
147 + state.apply_staged_writes()
148 +
149 + logits = torch.zeros((2, VOCAB_SIZE), device=DEVICE)
150 + out = _apply(
151 + state,
152 + logits,
153 + input_ids=[12, END_A],
154 + local_pos=[0, 1],
155 + )
156 +
157 + assert out[0, END_A] == pytest.approx(1.0e9)
158 + assert out[1, END_B] == pytest.approx(1.0e9)
159 +
160 +
161 +def test_v2_thinking_budget_stops_after_natural_end_marker():
162 + req_states = _make_req_states(
163 + [1, START, 10, END, 20, 21, 22],
164 + prompt_len=1,
165 + )
166 + state = ThinkingBudgetState(req_states, MockDistinctEndReasoningConfig())
167 + state.add_request(3, SamplingParams(thinking_token_budget=3))
168 + state.apply_staged_writes()
169 +
170 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
171 + out = _apply(state, logits, input_ids=[22], local_pos=[0])
172 +
173 + assert torch.all(out == 0)
174 +
175 +
176 +def test_v2_thinking_budget_ignores_plain_request():
177 + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
178 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
179 + state.add_request(3, SamplingParams())
180 + state.apply_staged_writes()
181 +
182 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
183 + out = _apply(state, logits, input_ids=[12], local_pos=[0])
184 +
185 + assert torch.all(out == 0)
186 +
187 +
188 +def test_v2_thinking_budget_latest_prefill_end_disables_forcing():
189 + req_states = _make_req_states(
190 + [1, START, 10, 11, 12, END, 13],
191 + prompt_len=1,
192 + )
193 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
194 + state.add_request(3, SamplingParams(thinking_token_budget=3))
195 + state.apply_staged_writes()
196 +
197 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
198 + out = _apply(state, logits, input_ids=[13], local_pos=[0])
199 +
200 + assert torch.all(out == 0)
201 +
202 +
203 +def test_v2_thinking_budget_uses_latest_prefill_start_boundary():
204 + req_states = _make_req_states(
205 + [1, START, 10, 11, 12, END, 13, START, 14, 15, 16],
206 + prompt_len=1,
207 + )
208 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
209 + state.add_request(3, SamplingParams(thinking_token_budget=3))
210 + state.apply_staged_writes()
211 +
212 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
213 + out = _apply(state, logits, input_ids=[16], local_pos=[0])
214 +
215 + assert out[0, END] == pytest.approx(1.0e9)
216 +
217 +
218 +def test_v2_thinking_budget_incrementally_scans_long_generation():
219 + """Guard against rescanning the full token history on every decode step."""
220 + tokens = [1, START, *([10] * 16382)]
221 + req_states = _make_req_states(tokens)
222 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
223 + state.add_request(3, SamplingParams(thinking_token_budget=32768))
224 + state.apply_staged_writes()
225 +
226 + _apply(state, torch.zeros((1, VOCAB_SIZE), device=DEVICE), [10], [0])
227 + assert state.cached_scan_pos[3].item() == len(tokens)
228 +
229 + req_states.all_token_ids.stage_write(3, len(tokens), [10])
230 + req_states.total_len.stage_write_elem(3, len(tokens) + 1)
231 + req_states.apply_staged_writes()
232 + _apply(state, torch.zeros((1, VOCAB_SIZE), device=DEVICE), [10], [0])
233 +
234 + assert state.cached_scan_pos[3].item() == len(tokens) + 1
235 +
236 +
237 +def test_v2_thinking_budget_clamps_oversized_budget():
238 + """Budgets beyond int32 must not crash and behave as unlimited."""
239 + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
240 + state = ThinkingBudgetState(req_states, MockReasoningConfig())
241 + state.add_request(3, SamplingParams(thinking_token_budget=2**40))
242 + state.apply_staged_writes()
243 +
244 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
245 + out = _apply(state, logits, input_ids=[12], local_pos=[0])
246 +
247 + assert torch.all(out == 0)
248 +
249 +
250 +def test_v2_thinking_budget_continues_end_prefix_from_prompt():
251 + """A resumed prompt ending with a partial forced-end marker must not
252 + restart the marker sequence and duplicate its first token."""
253 + req_states = _make_req_states([1, START, 10, 11, END_A], prompt_len=5)
254 + state = ThinkingBudgetState(req_states, MockMultiTokenEndReasoningConfig())
255 + state.add_request(3, SamplingParams(thinking_token_budget=3))
256 + state.apply_staged_writes()
257 +
258 + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
259 + out = _apply(state, logits, input_ids=[END_A], local_pos=[0])
260 +
261 + assert out[0, END_B] == pytest.approx(1.0e9)
262 + assert out[0, END_A] == 0
vllm/config/reasoning.py
+31 -13
@@ -13,10 +13,10 @@ from vllm.tokenizers import cached_tokenizer_from_config
13 class ReasoningConfig:
14 """Configuration for reasoning models.
15
16 - Set `reasoning_start_str` and `reasoning_end_str` to the strings that delimit
17 - the reasoning block (e.g. `"<think>"` and `"</think>"`). The
18 - corresponding token IDs are derived automatically via
19 - `initialize_token_ids` and are not intended to be set directly.
16 + Set `reasoning_start_str` and `reasoning_end_str` to the strings used to
17 + enter and forcibly terminate reasoning. The end string may include a
18 + transition phrase before the parser's natural reasoning end marker. Token
19 + IDs are derived automatically by `initialize_token_ids`.
20 """
21
22 reasoning_parser: str = ""
@@ -24,7 +24,7 @@ class ReasoningConfig:
24 reasoning_start_str: str = ""
25 """String that indicates the start of reasoning."""
26 reasoning_end_str: str = ""
27 - """String that indicates the end of reasoning content."""
27 + """String forced when the thinking budget is exhausted."""
28
29 _reasoning_start_token_ids: list[int] | None = field(
30 default=None, init=False, repr=False
@@ -34,8 +34,11 @@ class ReasoningConfig:
34 _reasoning_end_token_ids: list[int] | None = field(
35 default=None, init=False, repr=False
36 )
37 - """Private backing field for `reasoning_end_token_ids`. Set by
38 - `initialize_token_ids`. Not intended to be configured directly."""
37 + """Private backing field for forced reasoning end token IDs."""
38 + _natural_reasoning_end_token_ids: list[int] | None = field(
39 + default=None, init=False, repr=False
40 + )
41 + """Token IDs that naturally terminate reasoning, as defined by the parser."""
42
43 _enabled: bool = field(default=False, init=False, repr=False)
44 """Private field indicating whether reasoning token IDs have been initialized.
@@ -55,15 +58,20 @@ class ReasoningConfig:
58
59 @property
60 def reasoning_end_token_ids(self) -> list[int] | None:
58 - """Token IDs derived from `reasoning_end_str`. Set automatically by
59 - `initialize_token_ids`. Not intended to be configured directly."""
61 + """Token IDs forced when the thinking budget is exhausted."""
62 return self._reasoning_end_token_ids
63
64 + @property
65 + def natural_reasoning_end_token_ids(self) -> list[int] | None:
66 + """Token IDs that indicate the model naturally ended reasoning."""
67 + return self._natural_reasoning_end_token_ids
68 +
69 def initialize_token_ids(self, model_config: ModelConfig) -> None:
70 """Initialize reasoning token IDs from strings using the tokenizer."""
71 if (
72 self._reasoning_start_token_ids is not None
73 and self._reasoning_end_token_ids is not None
74 + and self._natural_reasoning_end_token_ids is not None
75 ):
76 self._enabled = True
77 return # Already initialized
@@ -71,9 +79,8 @@ class ReasoningConfig:
79 tokenizer = cached_tokenizer_from_config(model_config=model_config)
80 reasoning_start_str = self.reasoning_start_str
81 reasoning_end_str = self.reasoning_end_str
74 - if self.reasoning_parser is not None and (
75 - not reasoning_start_str or not reasoning_end_str
76 - ):
82 + natural_reasoning_end_str = ""
83 + if self.reasoning_parser:
84 parser_cls = ReasoningParserManager.get_reasoning_parser(
85 self.reasoning_parser
86 )
@@ -85,6 +92,10 @@ class ReasoningConfig:
92 end_token = reasoning_parser.reasoning_end_str
93 if end_token and not reasoning_end_str:
94 reasoning_end_str = end_token
95 + natural_reasoning_end_str = end_token or ""
96 +
97 + if not natural_reasoning_end_str:
98 + natural_reasoning_end_str = reasoning_end_str
99
100 if not reasoning_start_str or not reasoning_end_str:
101 # If we don't have valid strings to tokenize,
@@ -96,8 +107,15 @@ class ReasoningConfig:
107 self._reasoning_end_token_ids = tokenizer.encode(
108 reasoning_end_str, add_special_tokens=False
109 )
110 + self._natural_reasoning_end_token_ids = tokenizer.encode(
111 + natural_reasoning_end_str, add_special_tokens=False
112 + )
113
100 - if not self._reasoning_start_token_ids or not self._reasoning_end_token_ids:
114 + if (
115 + not self._reasoning_start_token_ids
116 + or not self._reasoning_end_token_ids
117 + or not self._natural_reasoning_end_token_ids
118 + ):
119 raise ValueError(
120 f"ReasoningConfig: failed to tokenize reasoning strings: "
121 f"reasoning_start_str='{self.reasoning_start_str}', "
vllm/config/vllm.py
-6
@@ -2377,12 +2377,6 @@ class VllmConfig:
2377 f"Model Runner V2 does not yet support: {', '.join(unsupported)}"
2378 )
2379
2380 - if self.reasoning_config is not None:
2381 - logger.warning_once(
2382 - "Model Runner V2 does not yet support the thinking_token_budget "
2383 - "request parameter. Set VLLM_USE_V2_MODEL_RUNNER=0 if this is required."
2384 - )
2385 -
2380 def validate_block_size(self) -> None:
2381 """Validate block_size against DCP and mamba constraints.
2382
vllm/v1/engine/input_processor.py
+9 -16
@@ -108,22 +108,15 @@ class InputProcessor:
108 self.tokenizer,
109 )
110
111 - if params.thinking_token_budget is not None:
112 - if (
113 - self.vllm_config.reasoning_config is None
114 - or not self.vllm_config.reasoning_config.enabled
115 - ):
116 - raise VLLMValidationError(
117 - "thinking_token_budget is set but reasoning_config is "
118 - "not configured. Please set --reasoning-parser "
119 - "and/or --reasoning-config to use thinking_token_budget."
120 - )
121 - if self.use_v2_model_runner:
122 - raise VLLMValidationError(
123 - "thinking_token_budget is not yet supported by the V2 "
124 - "model runner. Run vLLM with VLLM_USE_V2_MODEL_RUNNER=0 "
125 - "to use thinking_token_budget."
126 - )
111 + if params.thinking_token_budget is not None and (
112 + self.vllm_config.reasoning_config is None
113 + or not self.vllm_config.reasoning_config.enabled
114 + ):
115 + raise VLLMValidationError(
116 + "thinking_token_budget is set but reasoning_config is "
117 + "not configured. Please set --reasoning-parser "
118 + "and/or --reasoning-config to use thinking_token_budget."
119 + )
120 elif isinstance(params, PoolingParams):
121 supported_pooling_tasks = [
122 task for task in supported_tasks if task in POOLING_TASKS
vllm/v1/worker/gpu/model_runner.py
+1
@@ -371,6 +371,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
371 logprobs_mode=self.model_config.logprobs_mode,
372 num_speculative_tokens=self.decode_query_len,
373 use_fp64_gumbel=self.model_config.use_fp64_gumbel,
374 + reasoning_config=self.vllm_config.reasoning_config,
375 )
376 custom = self.model_state.custom_sampler(self.sampler)
377
vllm/v1/worker/gpu/sample/sampler.py
+22
@@ -6,6 +6,7 @@ import torch
6
7 import vllm.envs as envs
8 from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
9 +from vllm.config.reasoning import ReasoningConfig
10 from vllm.sampling_params import SamplingParams
11 from vllm.v1.sample.ops.topk_topp_sampler import (
12 apply_top_k_top_p,
@@ -24,6 +25,7 @@ from vllm.v1.worker.gpu.sample.logprob import (
25 from vllm.v1.worker.gpu.sample.output import SamplerOutput
26 from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
27 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
28 +from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
29 from vllm.v1.worker.gpu.states import RequestState
30
31
@@ -37,6 +39,7 @@ class Sampler:
39 logprobs_mode: LogprobsMode = "raw_logprobs",
40 num_speculative_tokens: int = 1,
41 use_fp64_gumbel: bool = False,
42 + reasoning_config: ReasoningConfig | None = None,
43 ):
44 self.logprobs_mode = logprobs_mode
45 self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
@@ -48,6 +51,7 @@ class Sampler:
51 self.logit_bias_state = LogitBiasState(max_num_reqs, device)
52 self.bad_words_state = BadWordsState(req_states)
53 self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
54 + self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
55 self.num_speculative_tokens = num_speculative_tokens
56 self.use_flashinfer = flashinfer_sampler_supported()
57
@@ -59,6 +63,7 @@ class Sampler:
63 self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)
64 self.bad_words_state.add_request(req_idx, sampling_params)
65 self.logprob_token_ids_state.add_request(req_idx, sampling_params)
66 + self.thinking_budget_state.add_request(req_idx, sampling_params)
67
68 def apply_staged_writes(self) -> None:
69 self.sampling_states.apply_staged_writes()
@@ -66,6 +71,7 @@ class Sampler:
71 self.logit_bias_state.apply_staged_writes()
72 self.bad_words_state.apply_staged_writes()
73 self.logprob_token_ids_state.apply_staged_writes()
74 + self.thinking_budget_state.apply_staged_writes()
75
76 def __call__(
77 self,
@@ -73,6 +79,7 @@ class Sampler:
79 input_batch: InputBatch,
80 ) -> SamplerOutput:
81 expanded_idx_mapping = input_batch.expanded_idx_mapping
82 + idx_mapping = input_batch.idx_mapping
83 idx_mapping_np = input_batch.idx_mapping_np
84 cu_num_logits_np = input_batch.cu_num_logits_np
85 expanded_local_pos = input_batch.expanded_local_pos
@@ -92,6 +99,7 @@ class Sampler:
99 sampled, processed_logits = self.sample(
100 logits,
101 expanded_idx_mapping,
102 + idx_mapping,
103 idx_mapping_np,
104 pos,
105 input_ids,
@@ -146,6 +154,7 @@ class Sampler:
154 self,
155 logits: torch.Tensor,
156 expanded_idx_mapping: torch.Tensor,
157 + idx_mapping: torch.Tensor,
158 idx_mapping_np: np.ndarray,
159 pos: torch.Tensor,
160 input_ids: torch.Tensor,
@@ -181,6 +190,17 @@ class Sampler:
190 expanded_local_pos,
191 )
192
193 + # Force the reasoning end marker once a request's thinking budget is
194 + # reached; applied before temperature so the forced token is always kept.
195 + self.thinking_budget_state.apply(
196 + logits,
197 + expanded_idx_mapping,
198 + idx_mapping,
199 + idx_mapping_np,
200 + input_ids,
201 + expanded_local_pos,
202 + )
203 +
204 # Apply temperature in place.
205 self.sampling_states.apply_temperature(
206 logits, expanded_idx_mapping, idx_mapping_np
@@ -219,6 +239,7 @@ class Sampler:
239 self,
240 logits: torch.Tensor,
241 expanded_idx_mapping: torch.Tensor,
242 + idx_mapping: torch.Tensor,
243 idx_mapping_np: np.ndarray,
244 pos: torch.Tensor,
245 input_ids: torch.Tensor,
@@ -228,6 +249,7 @@ class Sampler:
249 processed_logits = self.apply_sampling_params(
250 logits,
251 expanded_idx_mapping,
252 + idx_mapping,
253 idx_mapping_np,
254 pos,
255 input_ids,
vllm/v1/worker/gpu/sample/thinking_budget.py new
+428
@@ -0,0 +1,428 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import TYPE_CHECKING
4 +
5 +import numpy as np
6 +import torch
7 +
8 +from vllm.sampling_params import SamplingParams
9 +from vllm.triton_utils import tl, triton
10 +from vllm.utils.torch_utils import async_tensor_h2d
11 +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
12 +from vllm.v1.worker.gpu.states import RequestState
13 +
14 +if TYPE_CHECKING:
15 + from vllm.config.reasoning import ReasoningConfig
16 +
17 +_INT32_MAX = np.iinfo(np.int32).max
18 +_COLD_SCAN_BLOCK = 1024
19 +
20 +
21 +class ThinkingBudgetState:
22 + """Model Runner V2 state for per-request thinking token budgets."""
23 +
24 + def __init__(
25 + self,
26 + req_states: RequestState,
27 + reasoning_config: "ReasoningConfig | None",
28 + ):
29 + self.req_states = req_states
30 + self.max_num_reqs = req_states.max_num_reqs
31 + self.device = req_states.device
32 +
33 + start_ids = (
34 + []
35 + if reasoning_config is None
36 + else reasoning_config.reasoning_start_token_ids or []
37 + )
38 + end_ids = (
39 + []
40 + if reasoning_config is None
41 + else reasoning_config.reasoning_end_token_ids or []
42 + )
43 + natural_end_ids = (
44 + []
45 + if reasoning_config is None
46 + else reasoning_config.natural_reasoning_end_token_ids or []
47 + )
48 + self.enabled = bool(start_ids and end_ids and natural_end_ids)
49 + if not self.enabled:
50 + return
51 +
52 + self.thinking_token_budget = UvaBackedTensor(
53 + self.max_num_reqs, dtype=torch.int32
54 + )
55 + self.thinking_token_budget.np.fill(-1)
56 + self.thinking_token_budget.copy_to_uva()
57 + self.use_thinking_budget = np.zeros(self.max_num_reqs, dtype=bool)
58 +
59 + self.cached_last_start = torch.full(
60 + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device
61 + )
62 + self.cached_last_end = torch.full(
63 + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device
64 + )
65 + self.cached_scan_pos = torch.zeros(
66 + self.max_num_reqs, dtype=torch.int32, device=self.device
67 + )
68 + self._reset_reqs: list[int] = []
69 + self._budget_dirty = False
70 +
71 + self.reasoning_start_token_ids = torch.tensor(
72 + start_ids, dtype=torch.int32, device=self.device
73 + )
74 + self.natural_reasoning_end_token_ids = torch.tensor(
75 + natural_end_ids, dtype=torch.int32, device=self.device
76 + )
77 + self.reasoning_end_token_ids = torch.tensor(
78 + end_ids, dtype=torch.int32, device=self.device
79 + )
80 +
81 + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
82 + if not self.enabled:
83 + return
84 + budget = sampling_params.thinking_token_budget
85 + self.use_thinking_budget[req_idx] = budget is not None
86 + if budget is None:
87 + budget = -1
88 + else:
89 + budget = min(budget, _INT32_MAX)
90 + self._reset_reqs.append(req_idx)
91 + if self.thinking_token_budget.np[req_idx] != budget:
92 + self.thinking_token_budget.np[req_idx] = budget
93 + self._budget_dirty = True
94 +
95 + def apply_staged_writes(self) -> None:
96 + if not self.enabled:
97 + return
98 + if self._reset_reqs:
99 + idx = async_tensor_h2d(
100 + self._reset_reqs, dtype=torch.int64, device=self.device
101 + )
102 + self.cached_last_start.index_fill_(0, idx, -1)
103 + self.cached_last_end.index_fill_(0, idx, -1)
104 + self.cached_scan_pos.index_fill_(0, idx, 0)
105 + self._reset_reqs.clear()
106 + if self._budget_dirty:
107 + self.thinking_token_budget.copy_to_uva()
108 + self._budget_dirty = False
109 +
110 + def apply(
111 + self,
112 + logits: torch.Tensor,
113 + expanded_idx_mapping: torch.Tensor,
114 + idx_mapping: torch.Tensor,
115 + idx_mapping_np: np.ndarray,
116 + input_ids: torch.Tensor,
117 + expanded_local_pos: torch.Tensor,
118 + ) -> None:
119 + if not self.enabled or not np.any(self.use_thinking_budget[idx_mapping_np]):
120 + return
121 +
122 + apply_thinking_budget(
123 + logits,
124 + idx_mapping,
125 + expanded_idx_mapping,
126 + self.thinking_token_budget.gpu,
127 + self.req_states.all_token_ids.gpu,
128 + self.req_states.total_len.gpu,
129 + input_ids,
130 + expanded_local_pos,
131 + self.cached_last_start,
132 + self.cached_last_end,
133 + self.cached_scan_pos,
134 + self.reasoning_start_token_ids,
135 + self.natural_reasoning_end_token_ids,
136 + self.reasoning_end_token_ids,
137 + )
138 +
139 +
140 +@triton.jit
141 +def _load_effective_token(
142 + all_token_ids_ptr,
143 + all_token_ids_stride,
144 + input_ids_ptr,
145 + cur_req_first_pos,
146 + req_state_idx,
147 + total_len,
148 + pos,
149 +):
150 + if pos < total_len:
151 + return tl.load(all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos)
152 + # In decode/spec-decode, input_ids at local position 0 is the already
153 + # committed last sampled token. Effective draft-prefix positions start at
154 + # local position 1.
155 + input_pos = cur_req_first_pos + pos - total_len + 1
156 + return tl.load(input_ids_ptr + input_pos)
157 +
158 +
159 +@triton.jit
160 +def _update_committed_marker_cache_kernel(
161 + req_ids_ptr,
162 + thinking_token_budget_ptr,
163 + all_token_ids_ptr,
164 + all_token_ids_stride,
165 + total_len_ptr,
166 + cached_last_start_ptr,
167 + cached_last_end_ptr,
168 + cached_scan_pos_ptr,
169 + reasoning_start_token_ids_ptr,
170 + natural_reasoning_end_token_ids_ptr,
171 + START_LEN: tl.constexpr,
172 + NATURAL_END_LEN: tl.constexpr,
173 + MAX_LEN: tl.constexpr,
174 + BLOCK: tl.constexpr,
175 +):
176 + req_state_idx = tl.load(req_ids_ptr + tl.program_id(0))
177 + budget = tl.load(thinking_token_budget_ptr + req_state_idx)
178 + if budget < 0:
179 + return
180 +
181 + total_len = tl.load(total_len_ptr + req_state_idx)
182 + scan_pos = tl.load(cached_scan_pos_ptr + req_state_idx)
183 + last_start = tl.load(cached_last_start_ptr + req_state_idx)
184 + last_end = tl.load(cached_last_end_ptr + req_state_idx)
185 +
186 + if scan_pos > total_len:
187 + scan_pos = 0
188 + last_start = -1
189 + last_end = -1
190 +
191 + if scan_pos == 0 and last_start < 0 and last_end < 0:
192 + # Cold scan: walk backward in vectorized blocks, stopping at the first
193 + # block with a marker; only the relative order of the two positions
194 + # found matters below.
195 + block_hi = total_len
196 + while block_hi > 0 and last_start < 0 and last_end < 0:
197 + block_lo = block_hi - BLOCK
198 + if block_lo < 0:
199 + block_lo = 0
200 + offs = block_lo + tl.arange(0, BLOCK)
201 +
202 + start_match = (offs < block_hi) & (offs + START_LEN <= total_len)
203 + for j in tl.static_range(0, START_LEN):
204 + expected = tl.load(reasoning_start_token_ids_ptr + j)
205 + actual = tl.load(
206 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + offs + j,
207 + mask=offs + j < total_len,
208 + other=-1,
209 + )
210 + start_match = start_match & (actual == expected)
211 +
212 + end_match = (offs < block_hi) & (offs + NATURAL_END_LEN <= total_len)
213 + for j in tl.static_range(0, NATURAL_END_LEN):
214 + expected = tl.load(natural_reasoning_end_token_ids_ptr + j)
215 + actual = tl.load(
216 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + offs + j,
217 + mask=offs + j < total_len,
218 + other=-1,
219 + )
220 + end_match = end_match & (actual == expected)
221 +
222 + last_start = tl.max(tl.where(start_match, offs, -1), axis=0)
223 + last_end = tl.max(tl.where(end_match, offs, -1), axis=0)
224 + block_hi = block_lo
225 + else:
226 + for i in tl.range(scan_pos, total_len):
227 + if i + START_LEN <= total_len:
228 + start_match = True
229 + for j in tl.static_range(0, START_LEN):
230 + expected = tl.load(reasoning_start_token_ids_ptr + j)
231 + actual = tl.load(
232 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + i + j
233 + )
234 + start_match = start_match & (actual == expected)
235 + if start_match:
236 + last_start = i
237 +
238 + if i + NATURAL_END_LEN <= total_len:
239 + end_match = True
240 + for j in tl.static_range(0, NATURAL_END_LEN):
241 + expected = tl.load(natural_reasoning_end_token_ids_ptr + j)
242 + actual = tl.load(
243 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + i + j
244 + )
245 + end_match = end_match & (actual == expected)
246 + if end_match:
247 + last_end = i
248 +
249 + tl.store(cached_last_start_ptr + req_state_idx, last_start)
250 + tl.store(cached_last_end_ptr + req_state_idx, last_end)
251 + new_scan_pos = total_len - (MAX_LEN - 1)
252 + if new_scan_pos < 0:
253 + new_scan_pos = 0
254 + tl.store(cached_scan_pos_ptr + req_state_idx, new_scan_pos)
255 +
256 +
257 +@triton.jit
258 +def _thinking_budget_kernel(
259 + logits_ptr,
260 + logits_stride,
261 + expanded_idx_mapping_ptr,
262 + thinking_token_budget_ptr,
263 + all_token_ids_ptr,
264 + all_token_ids_stride,
265 + total_len_ptr,
266 + input_ids_ptr,
267 + expanded_local_pos_ptr,
268 + cached_last_start_ptr,
269 + cached_last_end_ptr,
270 + reasoning_start_token_ids_ptr,
271 + natural_reasoning_end_token_ids_ptr,
272 + reasoning_end_token_ids_ptr,
273 + START_LEN: tl.constexpr,
274 + NATURAL_END_LEN: tl.constexpr,
275 + END_LEN: tl.constexpr,
276 +):
277 + token_idx = tl.program_id(0).to(tl.int64)
278 + req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
279 + budget = tl.load(thinking_token_budget_ptr + req_state_idx)
280 + if budget < 0:
281 + return
282 +
283 + local_pos = tl.load(expanded_local_pos_ptr + token_idx)
284 + cur_req_first_pos = token_idx - local_pos
285 + total_len = tl.load(total_len_ptr + req_state_idx)
286 + effective_len = total_len + local_pos
287 +
288 + last_start = tl.load(cached_last_start_ptr + req_state_idx)
289 + last_end = tl.load(cached_last_end_ptr + req_state_idx)
290 +
291 + start_lo = total_len - START_LEN + 1
292 + if start_lo < 0:
293 + start_lo = 0
294 + for i in tl.range(start_lo, effective_len - START_LEN + 1):
295 + start_match = True
296 + for j in tl.static_range(0, START_LEN):
297 + expected = tl.load(reasoning_start_token_ids_ptr + j)
298 + actual = _load_effective_token(
299 + all_token_ids_ptr,
300 + all_token_ids_stride,
301 + input_ids_ptr,
302 + cur_req_first_pos,
303 + req_state_idx,
304 + total_len,
305 + i + j,
306 + )
307 + start_match = start_match & (actual == expected)
308 + if start_match:
309 + last_start = i
310 +
311 + end_lo = total_len - NATURAL_END_LEN + 1
312 + if end_lo < 0:
313 + end_lo = 0
314 + for i in tl.range(end_lo, effective_len - NATURAL_END_LEN + 1):
315 + end_match = True
316 + for j in tl.static_range(0, NATURAL_END_LEN):
317 + expected = tl.load(natural_reasoning_end_token_ids_ptr + j)
318 + actual = _load_effective_token(
319 + all_token_ids_ptr,
320 + all_token_ids_stride,
321 + input_ids_ptr,
322 + cur_req_first_pos,
323 + req_state_idx,
324 + total_len,
325 + i + j,
326 + )
327 + end_match = end_match & (actual == expected)
328 + if end_match:
329 + last_end = i
330 +
331 + if last_start < 0 or last_start <= last_end:
332 + return
333 +
334 + reasoning_start = last_start + START_LEN
335 + # If the request resumes from a prompt that already contains generated
336 + # reasoning content, count it against the remaining budget.
337 + num_reasoning_tokens = effective_len - reasoning_start
338 + if num_reasoning_tokens < budget:
339 + return
340 +
341 + # If the tail already ends with a prefix of the forced end sequence
342 + # (even from a resumed prompt), continue from the next marker token.
343 + end_prefix_len = 0
344 + max_prefix_len = END_LEN - 1
345 + if effective_len < max_prefix_len:
346 + max_prefix_len = effective_len
347 +
348 + for prefix_len in tl.static_range(1, END_LEN):
349 + if prefix_len <= max_prefix_len:
350 + prefix_match = True
351 + suffix_start = effective_len - prefix_len
352 + for j in tl.static_range(0, END_LEN):
353 + if j < prefix_len:
354 + expected = tl.load(reasoning_end_token_ids_ptr + j)
355 + actual = _load_effective_token(
356 + all_token_ids_ptr,
357 + all_token_ids_stride,
358 + input_ids_ptr,
359 + cur_req_first_pos,
360 + req_state_idx,
361 + total_len,
362 + suffix_start + j,
363 + )
364 + prefix_match = prefix_match & (actual == expected)
365 + if prefix_match:
366 + end_prefix_len = prefix_len
367 +
368 + force_token_id = tl.load(reasoning_end_token_ids_ptr + end_prefix_len)
369 + tl.store(logits_ptr + token_idx * logits_stride + force_token_id, 1.0e9)
370 +
371 +
372 +def apply_thinking_budget(
373 + logits: torch.Tensor,
374 + req_ids: torch.Tensor,
375 + expanded_idx_mapping: torch.Tensor,
376 + thinking_token_budget: torch.Tensor,
377 + all_token_ids: torch.Tensor,
378 + total_len: torch.Tensor,
379 + input_ids: torch.Tensor,
380 + expanded_local_pos: torch.Tensor,
381 + cached_last_start: torch.Tensor,
382 + cached_last_end: torch.Tensor,
383 + cached_scan_pos: torch.Tensor,
384 + reasoning_start_token_ids: torch.Tensor,
385 + natural_reasoning_end_token_ids: torch.Tensor,
386 + reasoning_end_token_ids: torch.Tensor,
387 +) -> None:
388 + num_tokens = logits.shape[0]
389 + start_len = reasoning_start_token_ids.shape[0]
390 + natural_end_len = natural_reasoning_end_token_ids.shape[0]
391 + end_len = reasoning_end_token_ids.shape[0]
392 +
393 + _update_committed_marker_cache_kernel[(req_ids.shape[0],)](
394 + req_ids,
395 + thinking_token_budget,
396 + all_token_ids,
397 + all_token_ids.stride(0),
398 + total_len,
399 + cached_last_start,
400 + cached_last_end,
401 + cached_scan_pos,
402 + reasoning_start_token_ids,
403 + natural_reasoning_end_token_ids,
404 + START_LEN=start_len,
405 + NATURAL_END_LEN=natural_end_len,
406 + MAX_LEN=max(start_len, natural_end_len),
407 + BLOCK=_COLD_SCAN_BLOCK,
408 + )
409 +
410 + _thinking_budget_kernel[(num_tokens,)](
411 + logits,
412 + logits.stride(0),
413 + expanded_idx_mapping,
414 + thinking_token_budget,
415 + all_token_ids,
416 + all_token_ids.stride(0),
417 + total_len,
418 + input_ids,
419 + expanded_local_pos,
420 + cached_last_start,
421 + cached_last_end,
422 + reasoning_start_token_ids,
423 + natural_reasoning_end_token_ids,
424 + reasoning_end_token_ids,
425 + START_LEN=start_len,
426 + NATURAL_END_LEN=natural_end_len,
427 + END_LEN=end_len,
428 + )
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+1
@@ -141,6 +141,7 @@ class RejectionSampler:
141 processed_logits = self.sampler.apply_sampling_params(
142 logits,
143 expanded_idx_mapping,
144 + idx_mapping,
145 idx_mapping_np,
146 pos,
147 draft_sampled,