[Bugfix][Reasoning] Fix thinking_token_budget not enforced on re-entry after forced end (#43757)
Signed-off-by: Ashwin Giridharan <girida@amazon.com> Signed-off-by: Cursor Agent <cursor-agent@cursor.com> Co-authored-by: Cursor Agent <cursor-agent@cursor.com> Co-authored-by: Simon Mo <simon.mo@hey.com>
Ashwin Giridharan committed
Jun 29, 2026 at 18:04 UTC
cda05ee8c4b39a3e8e8b74444bac11f12277d0a9
2 files changed
+146
-2
tests/v1/logits_processors/test_correctness.py
+124
@@ -1257,3 +1257,127 @@ def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index
1257
1258
assert h._state[0]["start_thinking"] == 0
1259
assert h._state[0]["end_thinking"] == expected_end_idx
1260
+
1261
+
1262
+# --- Thinking budget re-entry tests (issue #43708) ---
1263
+# Regression tests: after budget forces end-of-thinking token sequence,
1264
+# the state machine must detect and enforce budget on subsequent blocks.
1265
+
1266
+
1267
+class TestThinkingBudgetReentry:
1268
+ THINK_START = 100
1269
+ THINK_END_SINGLE = [200]
1270
+ THINK_END_MULTI = [200, 201, 202]
1271
+ BUDGET = 5
1272
+ CONTENT_TOKEN = 50
1273
+ THINK_TOKEN = 60
1274
+
1275
+ @staticmethod
1276
+ def _make_holder(end_token_ids: list[int]) -> ThinkingBudgetStateHolder:
1277
+ class FakeReasoningConfig:
1278
+ reasoning_start_token_ids = [TestThinkingBudgetReentry.THINK_START]
1279
+ reasoning_end_token_ids: list[int] = []
1280
+ enabled = True
1281
+
1282
+ cfg = FakeReasoningConfig()
1283
+ cfg.reasoning_end_token_ids = end_token_ids
1284
+ return ThinkingBudgetStateHolder(
1285
+ reasoning_config=cfg,
1286
+ max_num_seqs=8,
1287
+ num_spec_tokens=0,
1288
+ device=torch.device("cpu"),
1289
+ is_pin_memory=False,
1290
+ )
1291
+
1292
+ @staticmethod
1293
+ def _sync_batch(holder: ThinkingBudgetStateHolder, budget: int) -> None:
1294
+ holder.sync_batch(
1295
+ BatchUpdate(
1296
+ batch_size=1,
1297
+ removed=(),
1298
+ added=[(0, SamplingParams(thinking_token_budget=budget), None, [])],
1299
+ moved=(),
1300
+ )
1301
+ )
1302
+
1303
+ @staticmethod
1304
+ def _step(holder: ThinkingBudgetStateHolder, output_tok_ids: list[int]) -> None:
1305
+ holder.update_state(
1306
+ output_token_ids=[output_tok_ids],
1307
+ spec_token_ids=None,
1308
+ repeat_indices=None,
1309
+ )
1310
+
1311
+ def _exhaust_budget(self, holder: ThinkingBudgetStateHolder) -> list[int]:
1312
+ output = [self.THINK_START]
1313
+ self._step(holder, list(output))
1314
+ for _ in range(self.BUDGET):
1315
+ output.append(self.THINK_TOKEN)
1316
+ self._step(holder, list(output))
1317
+
1318
+ assert holder._state[0]["in_end"]
1319
+ return output
1320
+
1321
+ def _accept_end_tokens(
1322
+ self,
1323
+ holder: ThinkingBudgetStateHolder,
1324
+ output: list[int],
1325
+ end_token_ids: list[int],
1326
+ ) -> None:
1327
+ for tok in end_token_ids:
1328
+ output.append(tok)
1329
+ self._step(holder, list(output))
1330
+
1331
+ def test_single_token_end_reentry(self):
1332
+ holder = self._make_holder(self.THINK_END_SINGLE)
1333
+ self._sync_batch(holder, self.BUDGET)
1334
+
1335
+ output = self._exhaust_budget(holder)
1336
+ self._accept_end_tokens(holder, output, self.THINK_END_SINGLE)
1337
+
1338
+ for _ in range(3):
1339
+ output.append(self.CONTENT_TOKEN)
1340
+ self._step(holder, list(output))
1341
+
1342
+ output.append(self.THINK_START)
1343
+ self._step(holder, list(output))
1344
+ for _ in range(self.BUDGET):
1345
+ output.append(self.THINK_TOKEN)
1346
+ self._step(holder, list(output))
1347
+
1348
+ assert holder._state[0]["in_end"], (
1349
+ "Second thinking block must also be budget-enforced"
1350
+ )
1351
+
1352
+ def test_multi_token_end_reentry(self):
1353
+ holder = self._make_holder(self.THINK_END_MULTI)
1354
+ self._sync_batch(holder, self.BUDGET)
1355
+
1356
+ output = self._exhaust_budget(holder)
1357
+ self._accept_end_tokens(holder, output, self.THINK_END_MULTI)
1358
+
1359
+ assert not holder._state[0]["in_end"]
1360
+
1361
+ output.append(self.THINK_START)
1362
+ self._step(holder, list(output))
1363
+ for _ in range(self.BUDGET):
1364
+ output.append(self.THINK_TOKEN)
1365
+ self._step(holder, list(output))
1366
+
1367
+ assert holder._state[0]["in_end"], (
1368
+ "Immediate re-entry after multi-token end must be enforced"
1369
+ )
1370
+
1371
+ def test_single_block_not_broken(self):
1372
+ holder = self._make_holder(self.THINK_END_SINGLE)
1373
+ self._sync_batch(holder, self.BUDGET)
1374
+
1375
+ output = self._exhaust_budget(holder)
1376
+ self._accept_end_tokens(holder, output, self.THINK_END_SINGLE)
1377
+
1378
+ for _ in range(20):
1379
+ output.append(self.CONTENT_TOKEN)
1380
+ self._step(holder, list(output))
1381
+
1382
+ assert not holder._state[0]["in_end"]
1383
+ assert not holder._state[0]["in_think"]
vllm/v1/sample/thinking_budget_state.py
+22
-2
@@ -244,6 +244,7 @@ class ThinkingBudgetStateHolder:
244
"in_spec_mode": False,
245
"bonus_token_forced": False,
246
"continue_thinking": continue_thinking,
247
+ "scan_offset": 0,
248
}
249
250
def _update_think_state(self, state: dict[str, Any]) -> None:
@@ -258,20 +259,34 @@ class ThinkingBudgetStateHolder:
259
output_tok_ids = state.get("output_tok_ids", [])
260
if state["start_thinking"] == -1:
261
seq_len = len(self.think_start_token_ids)
262
+ scan_offset = state.get("scan_offset", 0)
263
start_thinking = self._find_last_sequence_index_from(
264
output_tok_ids,
265
self.think_start_token_ids,
264
- state["start_search_pos"] - (seq_len - 1),
266
+ max(scan_offset, state["start_search_pos"] - (seq_len - 1)),
267
)
268
+ if start_thinking >= 0 and scan_offset > 0:
269
+ # Re-entry after a forced end: budget was already exhausted
270
+ # in a prior block, so immediately force-close this one.
271
+ # scan_offset > 0 is only set after forced-end completion
272
+ # (never after natural end), so this won't block legitimate
273
+ # re-entries where budget remains.
274
+ state["start_thinking"] = start_thinking
275
+ state["in_think"] = False
276
+ state["in_end"] = True
277
+ state["end_count"] = 0
278
+ state["force_index"] = [0]
279
+ return
280
state["start_thinking"] = start_thinking
281
if start_thinking == -1:
282
state["start_search_pos"] = len(output_tok_ids)
283
if state["end_thinking"] == -1:
284
seq_len = len(self.think_end_token_ids)
285
+ scan_offset = state.get("scan_offset", 0)
286
end_thinking = self._find_last_sequence_index_from(
287
output_tok_ids,
288
self.think_end_token_ids,
274
- state["end_search_pos"] - (seq_len - 1),
289
+ max(scan_offset, state["end_search_pos"] - (seq_len - 1)),
290
)
291
state["end_thinking"] = end_thinking
292
if end_thinking == -1:
@@ -459,6 +474,11 @@ class ThinkingBudgetStateHolder:
474
"in_end": False,
475
"end_count": 0,
476
"check_count_down": state["thinking_token_budget"],
477
+ "start_thinking": -1,
478
+ "end_thinking": -1,
479
+ "think_count": 0,
480
+ "continue_thinking": False,
481
+ "scan_offset": len(state.get("output_tok_ids", [])),
482
}
483
)
484