执行摘要

• 一句话：修复自然思考结束后预算重入失效问题
• 推荐动作：建议阅读 thinking_budget_state.py 中 _update_think_state 的重写逻辑，理解 scan_offset 和早期检测的配合方式。测试代码 TestThinkingBudgetNaturalEndReentry 覆盖全面，可作为状态机边界测试的参考。

功能与动机

修复 issue #45974，补全 #43757 对强制结束重入的修复。根本原因是自然结束思考块时 start_thinking 和 end_thinking 未被重置，_find_last_sequence_index 从位置 0 搜索仍找到旧标记，导致新块被忽略。

实现拆解

1. 添加 scan_offset 字段：在状态条目中新增 scan_offset，初始化置 0；所有思考块退出（自然或强制）时重置 start_thinking/end_thinking 并设置 scan_offset 为当前输出长度，使后续搜索仅扫描新 token。
2. 早期自然结束检测：在预算倒计数的提前返回之前，检查 state["end_thinking"] > state["start_thinking"]，若成立则立即标记退出思考模式并重置位置；防止相邻 </think> 和 <think> 间无内容时倒计数不耗尽而错过重置点。
3. 简化搜索逻辑：删除独立的 _find_last_sequence_index_from 方法，统一使用 _find_last_sequence_index 在按 scan_offset 截取的子串上搜索；同时移除 start_search_pos/end_search_pos 字段以降低状态复杂度。
4. 统一退出路径：修改自然结束和强制结束两个分支，确保两种退出方式执行相同的位置标记重置和 scan_offset 推进。
5. 接口与测试配套：maybe_create_thinking_budget_state_holder 新增 is_pin_memory 参数，gpu_input_batch.py 传入 PIN_MEMORY；更新现有测试断言适配新行为。

关键文件：

• vllm/v1/sample/thinking_budget_state.py（模块 采样器；类别 source；类型 core-logic；符号 _update_think_state, _init_state_entry, _find_last_sequence_index_from）: 核心状态机修改：引入 scan_offset、早期自然结束检测，删除 _find_last_sequence_index_from 并简化搜索逻辑。是修复的主体。
• tests/v1/logits_processors/test_correctness.py（模块 测试；类别 test；类型 test-coverage；符号 TestThinkingBudgetNaturalEndReentry, _make_holder, FakeReasoningConfig, _sync_batch）: 新增 12 个测试用例覆盖自然/强制结束的各种重入场景，包括边缘情况（预算=1、中间无内容等），确保修复正确性。
• vllm/v1/worker/gpu_input_batch.py（模块 输入批处理；类别 source；类型 interface-adjustment）: 适配签名变更：调用 maybe_create_thinking_budget_state_holder 时传递 is_pin_memory 参数。

关键符号：ThinkingBudgetStateHolder._update_think_state, ThinkingBudgetStateHolder._init_state_entry, maybe_create_thinking_budget_state_holder, TestThinkingBudgetNaturalEndReentry._make_holder, TestThinkingBudgetNaturalEndReentry.test_natural_end_reentry_single_token

关键源码片段

vllm/v1/sample/thinking_budget_state.py

核心状态机修改：引入 scan_offset、早期自然结束检测，删除 _find_last_sequence_index_from 并简化搜索逻辑。是修复的主体。

# 在 ThinkingBudgetStateHolder._update_think_state 方法中新增的早期自然结束检测逻辑
if state["end_thinking"] > state["start_thinking"]:
    # 思考块已自然结束，重置状态以准备接收下一个可能的 <think> 块
    state["start_thinking"] = -1
    state["end_thinking"] = -1
    state["scan_offset"] = len(state.get("output_tok_ids", []))
    state["in_think"] = False
    state["think_count"] = 0

tests/v1/logits_processors/test_correctness.py

新增 12 个测试用例覆盖自然/强制结束的各种重入场景，包括边缘情况（预算=1、中间无内容等），确保修复正确性。

class TestThinkingBudgetNaturalEndReentry:
    """验证自然思考结束与预算重入的正确性。"""
​
    THINK_START = 100
    THINK_END_SINGLE = [200]
    BUDGET = 10
    THINK_TOKEN = 60
​
    @staticmethod
    def _make_holder(end_token_ids):
        """创建一个 ThinkingBudgetStateHolder 实例，使用虚假推理配置。"""
        from dataclasses import dataclass
        from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder
​
        @dataclass
        class FakeReasoningConfig:
            reasoning_start_token_ids: list[int]
            reasoning_end_token_ids: list[int]
            enabled: bool = True
​
        cfg = FakeReasoningConfig(
            reasoning_start_token_ids=[TestThinkingBudgetNaturalEndReentry.THINK_START],
            reasoning_end_token_ids=end_token_ids,
        )
        return ThinkingBudgetStateHolder(
            reasoning_config=cfg,
            max_num_seqs=8,
            num_spec_tokens=0,
            device=torch.device("cpu"),
            is_pin_memory=False,
        )
​
    def test_natural_end_reentry_single_token(self):
        """Block 1 自然结束，Block 2 被强制执行预算。"""
        holder = self._make_holder(self.THINK_END_SINGLE)
        self._sync_batch(holder, self.BUDGET)
        out = []
        # Block 1: 6 个思考 token + 自然结束 token [200]
        for _ in range(6):
            out.append(self.THINK_TOKEN)
        out.append(self.THINK_END_SINGLE[0])
        holder.update_state([out], None, None)
        # 验证状态机已重置并准备接收下一个 <think>
        assert not holder._state[0]["in_think"]
        assert holder._state[0]["scan_offset"] == len(out)

评论区精华

bbrowning 提出“思考预算追踪应通过 reasoning parsers 实现，因为它们是唯一可靠知道模型是否在 reasoning 的组件”。ashwing 回应称分离解析与预算的目的在于不同关注点：解析器识别何时开始/结束，持有者强制执行时长。simon-mo 询问是否会在 Q3 重新设计此部分，bbrowning 表示未积极进行但认为重新设计是唯一准确的方式。最终当前 PR 按原设计合并。

• 是否应通过 reasoning parsers 实现预算追踪 (design): 当前 PR 维持原有设计，但 simon-mo 询问是否在 Q3 重新设计此部分，bbrowning 表示尚未积极进行但认为重新设计是唯一准确的方式。

风险与影响

• 风险：核心风险为回归：状态机修改可能影响已有思考预算功能，特别是强制结束和单块场景。测试覆盖了 12 种组合（包括强制结束回归），但未覆盖与推测解码 (speculative decoding) 或流水线并行搭配的场景。此外，scan_offset 和早期检测逻辑增加分支复杂度，需关注后续维护。
• 影响：影响使用 thinking_token_budget 参数的用户，特别是 Qwen3 等带 reasoning parser 的模型。修复后，多轮对话中模型自然结束思考块后再次进入 <think> 时，预算能正确强制执行，避免无限推理。无 API/配置变更，无性能影响（仅 CPU 端状态逻辑）。
• 风险标记：核心状态机变更, 回归风险, 未覆盖推测解码组合

关联脉络

• PR #43757 [Bugfix][Reasoning] Fix thinking_token_budget not enforced on re-entry after forced end: 本 PR 是其自然结束场景的补充，两者共同解决 #43708 的完整问题。
• PR #43708 [Bug]: thinking_token_budget enforcement fails on multi-turn conversations: 原始 issue，描述了多轮对话中思考预算失效的问题，强制结束和自然结束两个修复共同解决。
• PR #45974 [Bug]: thinking_token_budget not enforced on re-entry after natural : 本 PR 修复的具体 issue，自然结束重入失败的详细复现。