[https://nvbugs/6485885][fix] Stop thinking-budget processor re-forcing the reasoning end tag#16785
Open
Wanli-Jiang wants to merge 1 commit into
Open
Conversation
…ng the reasoning end tag With the overlap scheduler, logits processors run one step behind the sampled tokens (request.get_tokens() does not yet include the previous iteration's token). ThinkingBudgetLogitsProcessor was stateless and re-derived its decision from that stale view, so the step right after forcing the reasoning end tag could not see it and forced it again. The raw stream then contained a doubled end tag (e.g. '</think></think>'); the reasoning parser consumes the first and the second leaks into message.content (reproduced 3/3 with nano-v3 + thinking_token_budget on nemotron-nano-3.5; clean with the overlap scheduler disabled). Track per-(request, beam) progress through the end sequence: continue from recorded progress instead of re-deriving from token_ids, stop once the sequence is complete, and reset when the view shows the block closed so later reasoning blocks are budgeted again. Also fixes multi-token end sequences restarting at the first token under the same lag. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
Walkthrough
ChangesReasoning-end progress handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unittest/llmapi/test_sampling_params.py (1)
211-261: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover request and beam isolation.
Both regressions use
req_id=0and a single beam, so they do not exercise the new(req_id, beam_idx)state boundary. Add cases where stale progress on one request and one beam cannot advance, suppress, or reset forcing for another.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_sampling_params.py` around lines 211 - 261, Extend the thinking-budget logits processor regression tests to cover state isolation by both req_id and beam_idx, using multiple requests and beams with stale token views. Verify progress from one (req_id, beam_idx) pair cannot advance, suppress, or reset end-sequence forcing for another pair, while preserving the existing stale-view behavior in both single- and multi-token end sequences.tensorrt_llm/llmapi/thinking_budget.py (1)
48-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd concrete annotations to the new Python interfaces.
tensorrt_llm/llmapi/thinking_budget.py#L48-L48: declare_end_progressasdict[tuple[int, int], int].tensorrt_llm/llmapi/thinking_budget.py#L65-L71: use built-in generic annotations and typekeyastuple[int, int].tests/unittest/llmapi/test_sampling_params.py#L204-L261: annotate_run()parameters/return type and test return types.Proposed fix
- self._end_progress: dict = {} + self._end_progress: dict[tuple[int, int], int] = {} - def _apply(self, req_id: int, token_ids: List[List[int]], logits: torch.Tensor) -> None: + def _apply(self, req_id: int, token_ids: list[list[int]], logits: torch.Tensor) -> None: - def _forced_token(self, key, token_ids: List[int]) -> Optional[int]: + def _forced_token( + self, key: tuple[int, int], token_ids: list[int] + ) -> Optional[int: - def _run(processor, token_ids): + def _run( + processor: ThinkingBudgetLogitsProcessor, + token_ids: list[list[int]], + ) -> torch.Tensor: - def test_thinking_budget_logits_processor_does_not_force_end_twice_on_stale_view(): + def test_thinking_budget_logits_processor_does_not_force_end_twice_on_stale_view() -> None: - def test_thinking_budget_logits_processor_continues_end_sequence_on_stale_view(): + def test_thinking_budget_logits_processor_continues_end_sequence_on_stale_view() -> None:As per coding guidelines, “Annotate every function” and “avoid
Any... prefer built-in generic types.” Based on learnings, Python 3.10+ syntax is supported.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/llmapi/thinking_budget.py` at line 48, The thinking-budget interfaces need concrete Python 3.10+ annotations. In tensorrt_llm/llmapi/thinking_budget.py lines 48-48, annotate _end_progress as dict[tuple[int, int], int]; in lines 65-71, use built-in generic annotations and type key as tuple[int, int]. In tests/unittest/llmapi/test_sampling_params.py lines 204-261, annotate _run() parameters and return type, along with the test functions’ return types, without using Any.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/llmapi/thinking_budget.py`:
- Around line 45-48: Update the request-finalization and cancellation flow
associated with _end_progress so all entries for the completed request’s
(req_id, beam_idx) keys are removed even when the reasoning end sequence is
never observed. Ensure state is cleared when a request finishes or is cancelled,
or make the processor request-scoped, while preserving normal end-sequence
cleanup.
---
Nitpick comments:
In `@tensorrt_llm/llmapi/thinking_budget.py`:
- Line 48: The thinking-budget interfaces need concrete Python 3.10+
annotations. In tensorrt_llm/llmapi/thinking_budget.py lines 48-48, annotate
_end_progress as dict[tuple[int, int], int]; in lines 65-71, use built-in
generic annotations and type key as tuple[int, int]. In
tests/unittest/llmapi/test_sampling_params.py lines 204-261, annotate _run()
parameters and return type, along with the test functions’ return types, without
using Any.
In `@tests/unittest/llmapi/test_sampling_params.py`:
- Around line 211-261: Extend the thinking-budget logits processor regression
tests to cover state isolation by both req_id and beam_idx, using multiple
requests and beams with stale token views. Verify progress from one (req_id,
beam_idx) pair cannot advance, suppress, or reset end-sequence forcing for
another pair, while preserving the existing stale-view behavior in both single-
and multi-token end sequences.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 74f86545-0917-46af-9a50-86a8a5d80245
📒 Files selected for processing (2)
tensorrt_llm/llmapi/thinking_budget.pytests/unittest/llmapi/test_sampling_params.py
Collaborator
|
PR_Github #61231 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61231 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61453 [ run ] triggered by Bot. Commit: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
With the overlap scheduler, logits processors run one step behind the sampled tokens (request.get_tokens() does not yet include the previous iteration's token). ThinkingBudgetLogitsProcessor was stateless and re-derived its decision from that stale view, so the step right after forcing the reasoning end tag could not see it and forced it again. The raw stream then contained a doubled end tag (e.g. ''); the reasoning parser consumes the first and the second leaks into message.content (reproduced 3/3 with nano-v3 + thinking_token_budget on nemotron-nano-3.5; clean with the overlap scheduler disabled).
Track per-(request, beam) progress through the end sequence: continue from recorded progress instead of re-deriving from token_ids, stop once the sequence is complete, and reset when the view shows the block closed so later reasoning blocks are budgeted again. Also fixes multi-token end sequences restarting at the first token under the same lag.
Dev Engineer Review
ThinkingBudgetLogitsProcessorto track per-request, per-beam reasoning-end progress, preventing duplicated end tags when the overlap scheduler exposes stale token views.client_id.QA Engineer Review
tests/unittest/llmapi/test_sampling_params.py.tests/integration/test_lists/.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.