[Perf][ThinkingBudget] reduce search space for thinking tokens (#46425)
Signed-off-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Walter Beller-Morales committed
Jun 24, 2026 at 11:02 UTC
7f99e80c3b69e11bb3a257a90aacbb2af182816c
2 files changed
+61
-4
tests/v1/logits_processors/test_correctness.py
+31
@@ -1226,3 +1226,34 @@ def test_thinking_budget_invalid_budget_rejected(invalid_budget):
1226
1227
with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
1228
SamplingParams(thinking_token_budget=invalid_budget)
1229
+
1230
+
1231
+def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index():
1232
+ """Test thinking budget enforced for a long thinking run,
1233
+ then a natural end marker."""
1234
+ h = ThinkingBudgetStateHolder(
1235
+ MockReasoningConfig(), 8, 0, torch.device("cpu"), False
1236
+ )
1237
+ h.sync_batch(
1238
+ BatchUpdate(
1239
+ batch_size=1,
1240
+ removed=(),
1241
+ added=[(0, SamplingParams(thinking_token_budget=10_000), None, [])],
1242
+ moved=(),
1243
+ )
1244
+ )
1245
+ start = MockReasoningConfig.reasoning_start_token_ids
1246
+ end = MockReasoningConfig.reasoning_end_token_ids
1247
+
1248
+ out: list[int] = list(start)
1249
+ h.update_state([out], None, None)
1250
+ for tok in range(500): # 500 filler thinking tokens, one decode step each
1251
+ out.append(tok)
1252
+ h.update_state([out], None, None)
1253
+ assert h._state[0]["end_thinking"] == -1 # not present yet
1254
+ expected_end_idx = len(out) # marker appended next
1255
+ out.extend(end)
1256
+ h.update_state([out], None, None)
1257
+
1258
+ assert h._state[0]["start_thinking"] == 0
1259
+ assert h._state[0]["end_thinking"] == expected_end_idx
vllm/v1/sample/thinking_budget_state.py
+30
-4
@@ -173,6 +173,19 @@ class ThinkingBudgetStateHolder:
173
return i
174
return -1
175
176
+ @staticmethod
177
+ def _find_last_sequence_index_from(
178
+ target_list: list[int], token_ids: list[int], search_start: int
179
+ ) -> int:
180
+ """Last occurrence of ``token_ids`` at or after ``search_start``."""
181
+ if not token_ids:
182
+ return -1
183
+ lo = max(0, search_start)
184
+ for i in range(len(target_list) - len(token_ids), lo - 1, -1):
185
+ if target_list[i : i + len(token_ids)] == token_ids:
186
+ return i
187
+ return -1
188
+
189
def _init_state_entry(
190
self, prompt_tok_ids: list[int] | None, thinking_token_budget: int
191
) -> dict[str, Any]:
@@ -226,6 +239,8 @@ class ThinkingBudgetStateHolder:
239
"force_index": [],
240
"start_thinking": start_thinking,
241
"end_thinking": -1,
242
+ "start_search_pos": 0,
243
+ "end_search_pos": 0,
244
"in_spec_mode": False,
245
"bonus_token_forced": False,
246
"continue_thinking": continue_thinking,
@@ -240,16 +255,27 @@ class ThinkingBudgetStateHolder:
255
state["force_index"] = []
256
return
257
258
+ output_tok_ids = state.get("output_tok_ids", [])
259
if state["start_thinking"] == -1:
244
- start_thinking = self._find_last_sequence_index(
245
- state.get("output_tok_ids", []), self.think_start_token_ids
260
+ seq_len = len(self.think_start_token_ids)
261
+ start_thinking = self._find_last_sequence_index_from(
262
+ output_tok_ids,
263
+ self.think_start_token_ids,
264
+ state["start_search_pos"] - (seq_len - 1),
265
)
266
state["start_thinking"] = start_thinking
267
+ if start_thinking == -1:
268
+ state["start_search_pos"] = len(output_tok_ids)
269
if state["end_thinking"] == -1:
249
- end_thinking = self._find_last_sequence_index(
250
- state.get("output_tok_ids", []), self.think_end_token_ids
270
+ seq_len = len(self.think_end_token_ids)
271
+ end_thinking = self._find_last_sequence_index_from(
272
+ output_tok_ids,
273
+ self.think_end_token_ids,
274
+ state["end_search_pos"] - (seq_len - 1),
275
)
276
state["end_thinking"] = end_thinking
277
+ if end_thinking == -1:
278
+ state["end_search_pos"] = len(output_tok_ids)
279
280
if state["start_thinking"] == -1:
281
return