[Bug]: thinking_token_budget not enforced on re-entry after natural </think>

[ashwing]

Describe the bug

When a model naturally ends a thinking block (emits </think> on its own before exhausting the budget), the ThinkingBudgetStateHolder state machine fails to track subsequent thinking blocks. A second <think> block in the same completion is never recognized as "in think" mode, so the budget is never enforced on it.

This means a model that naturally ends one thinking block early can then open a new <think> block and reason indefinitely with no budget enforcement.

Root Cause

After a natural </think>, start_thinking and end_thinking retain their values from Block 1 and are never reset. When Block 2's <think> appears, _find_last_sequence_index searches from position 0 (since scan_offset is not set for natural ends), finds the original Block 1 <think>, and since start_thinking (0) < end_thinking (7), hits the "exiting think mode" branch — so Block 2 is never recognized as entering think mode.

Compare with forced-end re-entry (fixed in #43757): scan_offset advances past Block 1 after the forced close completes (line 461), so the re-entry <think> is correctly detected as new.

Reproduction

import torch
from dataclasses import dataclass
from unittest.mock import MagicMock
from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder

THINK_START = 100
THINK_END = [200]
BUDGET = 10

@dataclass
class FakeReasoningConfig:
    reasoning_start_token_ids: list
    reasoning_end_token_ids: list
    enabled: bool = True

cfg = FakeReasoningConfig(
    reasoning_start_token_ids=[THINK_START],
    reasoning_end_token_ids=THINK_END,
)
holder = ThinkingBudgetStateHolder(
    reasoning_config=cfg, max_num_seqs=8,
    num_spec_tokens=0, device=torch.device("cpu"), is_pin_memory=False,
)

params = MagicMock()
params.thinking_token_budget = BUDGET
batch_update = MagicMock(removed=[], added=[(0, params, None, [])], moved=[])
holder.sync_batch(batch_update)

output = []

# Block 1: 6 tokens + natural </think>
output.append(THINK_START)
holder.update_state([list(output)], None, None)
for _ in range(6):
    output.append(60)  # think token
    holder.update_state([list(output)], None, None)
output.append(THINK_END[0])  # natural end
holder.update_state([list(output)], None, None)

# Content
for _ in range(3):
    output.append(50)
    holder.update_state([list(output)], None, None)

# Block 2: re-entry
output.append(THINK_START)
holder.update_state([list(output)], None, None)
for i in range(14):
    output.append(60)
    holder.update_state([list(output)], None, None)

state = holder._state[0]
assert state["in_end"], f"Block 2 should be budget-enforced after 14 tokens (budget={BUDGET}), but in_end={state['in_end']}"

Expected behavior

Block 2 should be budget-enforced. Either:

• Cumulative: tokens from Block 1 (6) count toward the total, so Block 2 gets cut at 4
• Per-block reset: Block 2 gets a fresh budget of 10, enforced after 10 tokens

Either policy is acceptable; what's not acceptable is zero enforcement.

Actual behavior

Block 2 never enters in_think=True, so the budget countdown never starts. The model can reason indefinitely in the second block.

Environment

• vLLM: vllm/vllm-openai:latest (also confirmed on current main)
• Affects all models using thinking_token_budget that can produce multiple think blocks in one completion

Related

• [Bug]: thinking_token_budget enforcement fails on multi-turn conversations when max_completion_tokens >> thinking_token_budget with ignore_eos:true #43708 — forced-end re-entry (fix in PR [Bugfix][Reasoning] Fix thinking_token_budget not enforced on re-entry after forced end #43757, scoped to scan_offset > 0 case)
• [Reasoning][Feature] Support for speculative decoding with thinking budget #34668 — original thinking budget implementation

[WangX0111]

Thanks for raising this!

I researched the related PR #43757 (which fixed the forced-end re-entry issue #43708) to understand how the thinking budget state machine handles re-entries.

Currently, the scan_offset update and state resets introduced in PR #43757 are scoped entirely to the forced-end completion path (END_COMPLETE). When a model naturally ends a thinking block via </think>, it bypasses this path. Consequently, scan_offset remains 0 and the cached state variables (start_thinking, end_thinking) are not reset, causing the second <think> block to be ignored by the scanner.

To resolve this cleanly without breaking the existing forced-end re-entry logic (where subsequent blocks should be immediately killed), I propose the following:

1. Introduce had_forced_end flag: Add a boolean "had_forced_end": False to the state dictionary initialized in _init_state_entry.
2. Identify Forced Ends: Set had_forced_end to True when the forced-end sequence completes (in END_COMPLETE).
3. Check Flag on Re-entry: In _update_think_state, check had_forced_end instead of scan_offset > 0 to decide if a new thinking block should be immediately force-closed.
4. Reset on Natural Ends: When a natural end is detected (absolute_end_pos >= 0), reset start_thinking to -1, end_thinking to -1, continue_thinking to False, and advance scan_offset to the current output length. Because had_forced_end remains False, the subsequent thinking blocks can be correctly scanned and budget-enforced (implementing the Per-block reset policy).

I have verified this approach locally with a standalone test case, and it resolves the natural end re-entry bug perfectly.

Since this issue is closely related to the unmerged PR #43757, I can submit a follow-up PR on top of #43757 once it gets merged.Please let me know if this direction looks good!

[ashwing]

thanks @WangX0111 for digging into this!

since i authored both #43757 and this issue, i've already got a fix implemented and tested that takes a slightly different approach — no had_forced_end flag needed.

the core insight: scan_offset alone is sufficient for both cases. rather than distinguishing forced vs natural ends, we just advance scan_offset and reset start_thinking/end_thinking on any think-block exit (natural or forced). this way _find_last_sequence_index is always scoped past completed blocks, and subsequent <think> tokens are detected fresh.

the approach:

1. scope _find_last_sequence_index to search from scan_offset (not from 0)
2. on natural end (start < end or only end found): reset markers + advance scan_offset
3. on forced-end completion: same reset (already in [Bugfix][Reasoning] Fix thinking_token_budget not enforced on re-entry after forced end #43757, extended here)

this also supersedes the forced-end fix in #43757 — same mechanism, broader scope. i'll open a unified PR shortly that covers both #43708 and #45974.

verified on Docker with/without the patch:

• without: 3/5 tests fail (both re-entry cases + forced end)
• with: all 5 pass (natural end single-token, multi-token, forced end, plus 2 no-regression tests)