archived internal review draft (c)

[valeriyischenko]

Fork-internal review PR. Base here is muse-glimmer-stacked so the diff shows only the kernel commit. The upstream PR will be this commit rebased onto vllm-project/vllm:main. The text below is the proposed upstream description.

Purpose

thinking_token_budget caps how many tokens a model may spend on reasoning. When the cap is hit, vLLM steers the model out of reasoning by injecting a configured "forced end" text into the output, one token per step. That works when the forced text is a single closing marker like </think>. It silently breaks when the forced text is longer and starts with the model's natural end marker. The kernel injects the first token, considers reasoning closed, and stops. The rest of the forced text never appears, and the model is not steered to the answer at all.

Forced text of that shape arises with message-framed models. For a model that wraps its reasoning in <think>...</think> tags, the closing tag is by itself a complete instruction to start answering, so the forced text and the natural end can be the same single marker. Muse-Glimmer (vllm-project#51655) frames reasoning as a message addressed to=self and closed by <|eom|>, and closing that message commits the model to nothing, because the natural next step can be another to=self message. A forced end for such a model has to both close the reasoning message and open the next channel, which makes it a multi-token sequence that necessarily starts with the natural end marker.

The MRv2 thinking-budget kernel (vllm/v1/worker/gpu/sample/thinking_budget.py) cannot deliver a forced reasoning-end sequence whose first tokens are the natural end marker. It forces exactly one token and silently drops the rest.

The kernel forces one token per decode step and re-derives its position in the forced sequence by matching against the tail of the output, but it only runs while a reasoning section is open:

if last_start < 0 or last_start <= last_end:
    return

If forced_ids[0] is the natural end, forcing it closes the section, so on the next step the guard sees a closed section and disengages. A config that asked for an N-token sequence gets 1 token, and the model is handed back control right after the natural end. In a message-framed model that leaves the model free to open another reasoning block with a fresh per-section budget.

The example configs in docs/features/reasoning_outputs.md are unaffected, which is why this has not been hit before. A transition phrase like ...I have the answer.</think> has its natural end at the end of the forced string, so the section stays open until the last forced token and tail-matching works.

Design

Resume the forced sequence from the natural end that closed the section rather than from the tail. The main risk in doing so is injecting the sequence twice, and two pieces of state prevent that:

1. cached_done_anchor remembers the position where a forced sequence has already been fully injected. The output can contain more end markers later, such as the <|eom|> that closes a second reasoning section or the <|eot|> that ends the answer. Without the anchor, each of those would look like a fresh budget overrun and the kernel would inject the sequence again.
2. When a request's cached state is reset and the kernel rescans the whole output from scratch, for example after the request is recomputed from its prompt, the rescan also looks for a forced sequence that is already complete in the output and restores the anchor. Without this, the rescan would see the overrun, find no anchor, and inject a second copy of a sequence the output already contains.

Speculative decoding can retract tokens after the anchor was set, so the anchor is revalidated against the actual output on every step. If the forced sequence is no longer intact because draft tokens were rejected, the kernel resumes forcing it instead of treating it as done. Whether a forced sequence extends its natural end is known from the config, so the whole path is a compile-time branch (CONTINUE_AFTER_NATURAL_END, a tl.constexpr). For every reasoning configuration that does not have this shape the flag is False, the branch is compiled out, and the generated kernel is identical to the current one.

Scope

This PR makes the kernel deliver the sequence the config asked for. It does not repair an observed production regression. On the checkpoint we tested (Muse-Glimmer-30B) the model produces the transition on its own after the forced natural end, so with no tools offered the unpatched and patched kernels give identical outputs (budgets 8/32/128, n=8 each). That behavior is a training-distribution property, not a guarantee.

The same experiments showed that a forced end string must never encode a routing decision. With one tool offered and the full transition <|eom|><|start|>assistant to=user<|message|> configured as the forced end, faithfully delivering it (this fix) sent 10 of 12 budgeted samples to the user instead of to the tool, and those answers fabricated the data the tool would have returned. A control run without a budget chose the tool in 6 of 6 samples. The unpatched kernel looked safer on that test only because it silently dropped 6 of the 7 configured tokens. The right configuration for such models is to force only the natural end and leave the recipient to the model. Issue vllm-project#44676 describes the same class of failure from the other direction. The docs change in this PR adds a note with this warning.

Delivering multi-token forced sequences faithfully is also the prerequisite for a safer future design in which only a channel-agnostic prefix (<|eom|><|start|>assistant to=) is forced and the single recipient position is constrained. That would guarantee reasoning ends while leaving the choice of the next channel to the model. The position-conditional constraint mechanism it needs does not exist yet and is out of scope here.

Related work

• [Feat] Support thinking_token_budget in Model Runner V2 vllm-project/vllm#46727 added this kernel and the natural/forced token-id split. It never exercises a forced string that begins with the natural end, which is the shape this PR fixes.
• [Reasoning] Break repeating reasoning loops by forcing the reasoning end sequence vllm-project/vllm#52677 forces reasoning_end_str to break repetition loops. It reads the same config field but triggers on repetition rather than budget exhaustion, has no continuation logic, and changes the V1 holder (thinking_budget_state.py), not the MRv2 kernel this PR fixes.
• [Bugfix][Core] Restore incremental thinking-budget marker scans vllm-project/vllm#52106 and fix(security): restore incremental marker search cursors in thinking budget state vllm-project/vllm#51133 fix a performance problem (a full rescan of the output on every step) in the V1 holder. They touch a different file and fix a different bug, so there is no overlap with this change.
• [Performance][MRV2] Cache logits-processing request state vllm-project/vllm#52329 fixed the sampler gate that decides whether logits processing runs at all. It is unrelated to the continuation logic, but it is what makes budget behavior observable on greedy requests in the first place.

Test Plan

pytest tests/v1/worker/test_gpu_thinking_budget.py -v

Four new tests over the smallest config with the broken shape (forced [END, END_A, END_B], natural [END]):

• continuation after the natural end
• continuation across a spec-decode row with rejected drafts
• a completed end is forced only once
• a new section that overruns the budget is still enforced

Test Result

All 17 tests in the file pass with the patch. With the unpatched kernel restored, exactly the two continuation tests fail and the two guard tests pass, which shows the tests exercise the new logic itself rather than the surrounding setup. The serving A/B on Muse-Glimmer-30B (only thinking_budget.py differing) is described in the Scope section. Without tools the outputs are identical. With tools it shows the hazard of a forced string that names a recipient, which the added docs note warns against.

---

AI assistance (Claude Code) was used in developing this change. All code was reviewed and validated by the submitter.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀