[https://nvbugs/6008710][fix] Adjust prompt logprobs to use the correct prompt token id by stnie · Pull Request #12499 · NVIDIA/TensorRT-LLM

stnie · 2026-03-24T11:13:55Z

Summary by CodeRabbit

Release Notes

Bug Fixes
- Strengthened validation for log probability calculations to ensure accuracy.
Tests
- Enhanced test coverage for log probability computations with improved validation checks.

Description

When determining the logprob of the prompt-token, the current implementation used the prompt token, which was used as input to create the logits/logprobs instead of the prompt-token which would come next. This PR adjusts this such that the prompt-token in the prompt-logprobs output are matched correctly

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)
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.

stnie · 2026-03-24T11:14:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-24T11:20:25Z

PR_Github #40112 [ run ] triggered by Bot. Commit: 5a33ac2 Link to invocation

tensorrt-cicd · 2026-03-24T21:16:22Z

PR_Github #40112 [ run ] completed with state SUCCESS. Commit: 5a33ac2
/LLM/main/L0_MergeRequest_PR pipeline #31261 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

stnie · 2026-03-25T08:23:02Z

/bot run

tensorrt-cicd · 2026-03-25T08:29:11Z

PR_Github #40295 [ run ] triggered by Bot. Commit: 5a33ac2 Link to invocation

tensorrt-cicd · 2026-03-25T16:44:16Z

PR_Github #40295 [ run ] completed with state SUCCESS. Commit: 5a33ac2
/LLM/main/L0_MergeRequest_PR pipeline #31406 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

stnie · 2026-03-30T17:03:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-30T17:09:19Z

PR_Github #40752 [ run ] triggered by Bot. Commit: 8940d49 Link to invocation

tensorrt-cicd · 2026-03-31T03:10:03Z

PR_Github #40752 [ run ] completed with state SUCCESS. Commit: 8940d49
/LLM/main/L0_MergeRequest_PR pipeline #31770 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

stnie · 2026-03-31T08:14:47Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-31T08:27:57Z

PR_Github #40909 [ run ] triggered by Bot. Commit: 8940d49 Link to invocation

stnie · 2026-03-31T15:16:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-31T15:22:30Z

PR_Github #40963 [ run ] triggered by Bot. Commit: 1b3b8cf Link to invocation

tensorrt-cicd · 2026-03-31T15:22:45Z

PR_Github #40909 [ run ] completed with state ABORTED. Commit: 8940d49
LLM/main/L0_MergeRequest_PR #31910 (Blue Ocean) completed with status: ABORTED

Link to invocation

tensorrt-cicd · 2026-03-31T23:48:49Z

PR_Github #40963 [ run ] completed with state FAILURE. Commit: 1b3b8cf
/LLM/main/L0_MergeRequest_PR pipeline #31949 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

Signed-off-by: Stefan Niebler <[email protected]>

… accurate computation - Adjusted the prompt token ID handling in `_compute_pytorch_prompt_logprobs` to append the first generated token ID. - Modified unit tests to reflect changes in prompt logprobs calculations, ensuring consistency with the updated logic. Signed-off-by: Stefan Niebler <[email protected]>

- replace key,value with token_id, logprob_obj - replace cannot with must not Signed-off-by: Stefan Niebler <[email protected]>

stnie · 2026-04-01T12:34:42Z

/bot run --disable-fail-fast

coderabbitai · 2026-04-01T12:40:15Z

📝 Walkthrough

Walkthrough

Changes modify how prompt logprobs are computed and validated in the executor and test suite. The executor now reconstructs prompt token IDs by combining the first generated token with remaining prompt tokens and adds validation assertions. The test suite extends validation logic with offset parameters and cross-validation between prompt and generation logprobs.

Changes

Cohort / File(s)	Summary
Executor prompt logprobs computation `tensorrt_llm/executor/base_worker.py`	Modified `_compute_pytorch_prompt_logprobs` to tighten assertions on `context_logits` and validate non-`None` results. Changed token ID construction to derive from `result.output_token_ids[0][:1]` prepended to the remaining prompt tokens, replacing direct use of the full `prompt_token_ids` list. Added assertion validating extracted first generation token is non-empty.
Test suite validation logic `tests/unittest/_torch/sampler/test_logits_logprobs.py`	Extended `check_logprobs` helper with `logprobs_offset` parameter for token-id indexing alignment. Updated token ID assembly to concatenate `output.prompt_token_ids[1:]` with first generated token. Added context token checking using the full prompt plus first generated token. Introduced cross-validation block comparing final context logprobs with first generation logprobs when both `prompt_logprobs_k` and `logprobs_k` are set, including assertions on sampled token consistency and top-n token logprob alignment.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	❓ Inconclusive	The description explains the issue and solution, but the 'Test Coverage' section is empty and the PR checklist lacks specific test references despite significant test file modifications.	Fill in the 'Test Coverage' section with specific test case names (e.g., test_sampled_token_always_in_prompt_logprobs) that validate the prompt logprobs corrections.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check	✅ Passed	The title clearly and specifically summarizes the main change: fixing prompt logprobs to use the correct prompt token id, which aligns with the PR objectives and code changes.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

tensorrt-cicd · 2026-04-01T12:40:16Z

PR_Github #41197 [ run ] triggered by Bot. Commit: 82eda0d Link to invocation

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/executor/base_worker.py (1)

858-866: Token alignment fix looks correct; consider guarding against empty output_token_ids.

The logic correctly aligns prompt token IDs with context logits by offsetting by 1 and appending the first generated token. However, line 862 accesses result.output_token_ids[0] which could raise IndexError if output_token_ids is an empty list, before the assertion on line 863 can catch it.

If you want the assertion message to be shown instead of a raw IndexError:

🛡️ Optional defensive check

     result = response.result.get_result()
     assert result is not None, "result must not be None when prompt_logprobs is requested."
+    assert result.output_token_ids and result.output_token_ids[0], \
+        "output_token_ids must not be empty when prompt_logprobs is requested."
     # Single element list
-    first_generation_token = result.output_token_ids[0][:1]
-    assert first_generation_token, "first generation token must not be empty when prompt_logprobs is requested."
+    first_generation_token = result.output_token_ids[0][:1]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/executor/base_worker.py` around lines 858 - 866, Before indexing
result.output_token_ids in base_worker.py, add a defensive assertion that
result.output_token_ids is non-empty (e.g., assert result.output_token_ids,
"result.output_token_ids must not be empty when prompt_logprobs is requested.")
so the assertion fires instead of an IndexError; then safely compute
first_generation_token from result.output_token_ids[0][:1] and continue building
prompt_token_ids from generation_result._generation_request.prompt_token_ids[1:]
+ first_generation_token, keeping the existing context_logits checks intact.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tensorrt_llm/executor/base_worker.py`:
- Around line 858-866: Before indexing result.output_token_ids in
base_worker.py, add a defensive assertion that result.output_token_ids is
non-empty (e.g., assert result.output_token_ids, "result.output_token_ids must
not be empty when prompt_logprobs is requested.") so the assertion fires instead
of an IndexError; then safely compute first_generation_token from
result.output_token_ids[0][:1] and continue building prompt_token_ids from
generation_result._generation_request.prompt_token_ids[1:] +
first_generation_token, keeping the existing context_logits checks intact.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a21726a2-34b9-46d4-bb73-122fa1885475

📥 Commits

Reviewing files that changed from the base of the PR and between 3a71062 and 82eda0d.

📒 Files selected for processing (2)

tensorrt_llm/executor/base_worker.py
tests/unittest/_torch/sampler/test_logits_logprobs.py

tensorrt-cicd · 2026-04-01T22:53:59Z

PR_Github #41197 [ run ] completed with state SUCCESS. Commit: 82eda0d
/LLM/main/L0_MergeRequest_PR pipeline #32158 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

stnie · 2026-04-02T06:40:28Z

/bot run

tensorrt-cicd · 2026-04-02T06:46:00Z

PR_Github #41364 [ run ] triggered by Bot. Commit: 82eda0d Link to invocation

tensorrt-cicd · 2026-04-02T12:17:00Z

PR_Github #41364 [ run ] completed with state SUCCESS. Commit: 82eda0d
/LLM/main/L0_MergeRequest_PR pipeline #32307 completed with status: 'SUCCESS'

CI Report

Link to invocation

…ct prompt token id (NVIDIA#12499) Signed-off-by: Stefan Niebler <[email protected]>

github-actions Bot assigned stnie Mar 24, 2026

stnie requested review from Funatiq and venkywonka March 24, 2026 11:14

Funatiq approved these changes Mar 24, 2026

View reviewed changes

stnie force-pushed the develop/sampler/fix_prompt_logprobs branch 2 times, most recently from bf1dd15 to 8940d49 Compare March 30, 2026 17:02

stnie requested a review from Funatiq March 30, 2026 17:03

Funatiq approved these changes Mar 31, 2026

View reviewed changes

Comment thread tensorrt_llm/executor/base_worker.py Outdated

Comment thread tensorrt_llm/executor/base_worker.py Outdated

Comment thread tests/unittest/_torch/sampler/test_logits_logprobs.py Outdated

stnie force-pushed the develop/sampler/fix_prompt_logprobs branch from 8940d49 to 4e40ea1 Compare March 31, 2026 15:13

stnie added 3 commits April 1, 2026 14:34

[None][fix] Adjust prompt logprobs to use the correct prompt token id

d6093c3

Signed-off-by: Stefan Niebler <[email protected]>

address review comments P1

82eda0d

- replace key,value with token_id, logprob_obj - replace cannot with must not Signed-off-by: Stefan Niebler <[email protected]>

stnie force-pushed the develop/sampler/fix_prompt_logprobs branch from 1b3b8cf to 82eda0d Compare April 1, 2026 12:34

stnie marked this pull request as ready for review April 1, 2026 12:34

stnie requested a review from a team as a code owner April 1, 2026 12:34

coderabbitai Bot reviewed Apr 1, 2026

View reviewed changes

stnie changed the title ~~[None][fix] Adjust prompt logprobs to use the correct prompt token id~~ [https:://nvbugs/6008710][fix] Adjust prompt logprobs to use the correct prompt token id Apr 2, 2026

stnie changed the title ~~[https:://nvbugs/6008710][fix] Adjust prompt logprobs to use the correct prompt token id~~ [https://nvbugs/6008710][fix] Adjust prompt logprobs to use the correct prompt token id Apr 2, 2026

stnie enabled auto-merge (squash) April 2, 2026 15:36

venkywonka approved these changes Apr 2, 2026

View reviewed changes

stnie merged commit d0ac8fc into NVIDIA:main Apr 2, 2026
10 of 11 checks passed

Conversation

stnie commented Mar 24, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Release Notes

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

stnie commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

stnie commented Mar 25, 2026

Uh oh!

tensorrt-cicd commented Mar 25, 2026

Uh oh!

tensorrt-cicd commented Mar 25, 2026

Uh oh!

stnie commented Mar 30, 2026

Uh oh!

tensorrt-cicd commented Mar 30, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

stnie commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

stnie commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

stnie commented Apr 1, 2026

Uh oh!

coderabbitai Bot commented Apr 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 inconclusive)

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

stnie commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

stnie commented Mar 24, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Apr 1, 2026 •

edited

Loading