Skip to content

fix(vllm): record logprob for the sampled token - #3297

Merged
terrykong merged 1 commit into
NVIDIA-NeMo:mainfrom
HollowMan6:log_probs
Jul 20, 2026
Merged

fix(vllm): record logprob for the sampled token#3297
terrykong merged 1 commit into
NVIDIA-NeMo:mainfrom
HollowMan6:log_probs

Conversation

@HollowMan6

Copy link
Copy Markdown
Contributor

What does this PR do ?

Record logprob for the sampled token

Issues

#2608 (comment)

Fix vLLM generation output processing to record the log probability of the token actually sampled, rather than relying on the first entry in the per-position logprob dictionary.

vLLM returns each generated position as a mapping from token ID to logprob metadata. When that mapping contains multiple tokens, the first entry is not guaranteed to be the sampled token.

For example, if token 880 was sampled:

{
    512: Logprob(logprob=-0.10, rank=1),
    880: Logprob(logprob=-2.30, rank=2),
}

The previous implementation recorded -0.10 by taking the first dictionary entry. The correct value for the generated token is -2.30.

So Pair each generated token ID with its corresponding logprob dictionary and look up the logprob by that token ID.

The behavior is unchanged when vLLM returns only the sampled token.

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

Fix vLLM generation output processing to record the log probability of the
token actually sampled, rather than relying on the first entry in the
per-position logprob dictionary.

vLLM returns each generated position as a mapping from token ID to logprob
metadata. When that mapping contains multiple tokens, the first entry is not
guaranteed to be the sampled token.

For example, if token 880 was sampled:

{
    512: Logprob(logprob=-0.10, rank=1),
    880: Logprob(logprob=-2.30, rank=2),
}

The previous implementation recorded -0.10 by taking the first dictionary
entry. The correct value for the generated token is -2.30.

So Pair each generated token ID with its corresponding logprob dictionary and
look up the logprob by that token ID.

The behavior is unchanged when vLLM returns only the sampled token.

Signed-off-by: Hollow Man <hollowman@opensuse.org>
Copilot AI review requested due to automatic review settings July 20, 2026 20:41
@HollowMan6
HollowMan6 requested review from a team as code owners July 20, 2026 20:41
@copy-pr-bot

copy-pr-bot Bot commented Jul 20, 2026

Copy link
Copy Markdown

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@HollowMan6 HollowMan6 added bug Something isn't working CI:L0 Run doctests and unit tests labels Jul 20, 2026
@NVIDIA-NeMo NVIDIA-NeMo deleted a comment from cursor Bot Jul 20, 2026
@HollowMan6 HollowMan6 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L0 Run doctests and unit tests labels Jul 20, 2026
@HollowMan6

Copy link
Copy Markdown
Contributor Author

/ok to test bdb414e

@terrykong
terrykong merged commit 78965b3 into NVIDIA-NeMo:main Jul 20, 2026
65 checks passed
@HollowMan6
HollowMan6 deleted the log_probs branch July 20, 2026 22:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working CI:L1 Run doctests, unit tests, and functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants