[Bug]: thinking_token_budget restarts a multi-token reasoning end tag the model already began

System Info

No GPU needed. ThinkingBudgetLogitsProcessor is pure Python plus torch, so this reproduces on CPU.

• TensorRT-LLM branch: main
• Commit: 395985c025c8d1cf5aa842bc752b337ba88721b6
• Python 3.12, torch 2.14.0

Who can help?

Whoever owns tensorrt_llm/llmapi/thinking_budget.py (the file came in with #14665 and was last fixed by #16785).

Description

When the reasoning end string tokenizes to more than one id, and the model starts emitting that end sequence on its own just below the budget, the processor forces the first end token again instead of continuing from where the model left off. The stream then carries a doubled end tag prefix such as </ </ think>, the reasoning parser never matches a clean </think>, and reasoning text leaks into message.content.

_forced_token has two budget checks that disagree with each other. The first one discounts the end tokens the model already emitted:

partial_end_len = _longest_suffix_prefix_len(token_ids, self.reasoning_end_token_ids)
if (partial_end_len > 0
        and reasoning_token_count - partial_end_len >= self.thinking_token_budget):
    self._end_progress[key] = partial_end_len + 1
    return self.reasoning_end_token_ids[partial_end_len]

The fallback right below it counts those same tokens as reasoning content:

if reasoning_token_count >= self.thinking_token_budget:
    self._end_progress[key] = 1
    return self.reasoning_end_token_ids[0]

So in the band budget - partial_end_len <= reasoning_token_count < budget the first branch correctly declines (content is still under budget) and the fallback fires anyway, restarting the end sequence at index 0.

This is the same class of failure that #16785 fixed for the stale-token-view path. That PR added the _end_progress bookkeeping but left the budget arithmetic alone, and its tests cover content == budget and content well under budget, not the band in between.

Reproduction

import torch
from tensorrt_llm.llmapi.thinking_budget import ThinkingBudgetLogitsProcessor

p = ThinkingBudgetLogitsProcessor(
    thinking_token_budget=3,
    reasoning_start_token_ids=[1],
    reasoning_end_token_ids=[2, 3],   # e.g. '</' + 'think>'
)
logits = torch.zeros(1, 1, 8)
p(0, logits, [[1, 5, 6, 2]], None, None)   # content is 5, 6, so 2 tokens, under budget
print(int(logits[0, 0].argmax()))          # prints 2; nothing should be forced yet

Decoding greedily on from there, with a model that wants to finish the tag it started:

main:     [1, 5, 6, 2, 2, 3, 7]     <think> ... </  </think>
expected: [1, 5, 6, 2, 3, 7, 7]     <think> ... </think>

Expected behavior

End tokens the model produced itself are not reasoning content, so they should not count against the budget, and when the budget is spent the processor should resume at reasoning_end_token_ids[partial_end_len] instead of restarting at index 0.

Scope

This only bites when the end sequence is more than one token, which is the normal case for any parser whose tag is not a single added token in the tokenizer. On the Qwen2.5 tokenizer, </mm:think> (MiniMax parser) encodes to 5 ids and <channel|> (harmony) to 4. Models that carry </think> as one added token, such as DeepSeek-R1-Distill, are not affected.

I have a fix and a regression test ready and will open a PR against this issue.

Start in tensorrt_llm/llmapi/thinking_budget.py, focusing on ThinkingBudgetLogitsProcessor._forced_token and the two budget checks. Reproduce the supplied CPU example, then add a regression test for a partially emitted multi-token end sequence. Done means the processor does not restart the end tag and reasoning text does not leak into message.content.