[TRTLLM-9735][feat] Add processed logprobs functionality to TorchSampler #9675

stnie · 2025-12-03T17:21:40Z

Summary by CodeRabbit

Release Notes

New Features
- Added support for configurable log probability modes ("raw" or "processed") to control logprob output format.
- Enhanced per-token sampled log probability tracking with top-k support.
- Improved sampling strategy handling with temperature state management.
Bug Fixes
- Fixed log probability parameter handling to respect explicit false values.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Description

Iintroduces a new optional parameter, logprobs_mode, to the SamplingParams and LlmRequest classes, allowing users to specify the mode of log probabilities to return.
Create process_logprobs function to remove logprobs processing code from process_requests.
add batching based on logprobs_mode to sample_batched_by_strategy
additionally return processed logits from sampling

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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

stnie · 2025-12-03T17:44:55Z

/bot run

tensorrt-cicd · 2025-12-03T17:51:33Z

PR_Github #26827 [ run ] triggered by Bot. Commit: 0b21e18

tensorrt-cicd · 2025-12-03T19:16:20Z

PR_Github #26827 [ run ] completed with state SUCCESS. Commit: 0b21e18
/LLM/main/L0_MergeRequest_PR pipeline #20432 completed with status: 'FAILURE'

stnie · 2025-12-05T16:56:22Z

/bot run

tensorrt-cicd · 2025-12-05T17:02:20Z

PR_Github #27140 [ run ] triggered by Bot. Commit: d64119a

tensorrt-cicd · 2025-12-05T18:49:28Z

PR_Github #27140 [ run ] completed with state SUCCESS. Commit: d64119a
/LLM/main/L0_MergeRequest_PR pipeline #20709 completed with status: 'FAILURE'

tensorrt_llm/sampling_params.py

tests/unittest/_torch/sampler/test_logits_logprobs.py

stnie · 2025-12-08T11:04:25Z

/bot run --disable-fail-fast

tensorrt-cicd · 2025-12-08T11:10:06Z

PR_Github #27296 [ run ] triggered by Bot. Commit: bfb2195

ixlmar

Reviewed code minus tests.

Re the type annotations, I recommend checking using some type checker (e.g. in the IDE) as long as we don't have that in the CI. Ideally, we could maintain type correctness for sampling_utils.py, flashinfer_sampling_utils.py, and the newer parts of sampler.py.

tensorrt_llm/_torch/pyexecutor/llm_request.py

tensorrt_llm/_torch/pyexecutor/sampler.py

tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py

tensorrt_llm/executor/result.py

tensorrt-cicd · 2025-12-08T18:25:39Z

PR_Github #27296 [ run ] completed with state SUCCESS. Commit: bfb2195
/LLM/main/L0_MergeRequest_PR pipeline #20847 completed with status: 'FAILURE'

stnie · 2025-12-09T15:52:11Z

/bot run

tensorrt-cicd · 2025-12-09T15:57:51Z

PR_Github #27528 [ run ] triggered by Bot. Commit: 18531f1

tensorrt-cicd · 2025-12-09T17:38:44Z

PR_Github #27528 [ run ] completed with state SUCCESS. Commit: 18531f1
/LLM/main/L0_MergeRequest_PR pipeline #21006 completed with status: 'FAILURE'

pyproject.toml

tensorrt_llm/_torch/pyexecutor/sampler.py

stnie · 2025-12-12T17:16:53Z

/bot run --disable-fail-fast

tensorrt_llm/sampling_params.py

litaotju · 2025-12-31T06:38:41Z

tensorrt_llm/executor/result.py

+            if tokens is not None:
+                for t in range(logprobs.size(0)):
+                    token_id = tokens[t]
+                    token_logprob = logprobs[t, token_id].item()


I think this will cause a device sync. Is this really what we want here? Does this function sync before this PR?

I think this is for feature completeness for TRT flow and is not perf tuned. But I think we are considering deprecating the TRT flow so the perf tunning priority is low?
@hchings for confirm.

tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py

litaotju · 2025-12-31T06:45:47Z

Could you pls run some performance tests and make sure there is no addtional sync in ths sampler part? Which can easily introduce regressions based on recent changes on sampler.

tensorrt_llm/_torch/pyexecutor/llm_request.py

tensorrt_llm/_torch/pyexecutor/py_executor.py

tensorrt_llm/_torch/pyexecutor/llm_request.py

tensorrt_llm/_torch/pyexecutor/sampler.py

litaotju · 2025-12-31T06:59:06Z

tensorrt_llm/_torch/pyexecutor/sampler.py

-    return resolve_sampling_strategy(params, vocab_size=vocab_size)
+    if not hasattr(request, "py_sampling_strategy") or _get_max_beam_width(request) > 1:
+        params = _request_get_sampling_params(request)
+        request.py_sampling_strategy = resolve_sampling_strategy(params, vocab_size=vocab_size)


I noticed this function changes its previous behavior that don't modify the input request to now modifying it.

Is this really what we want here?

I think this is to avoid repeated computation by only compute and store on request.py_sampling_strategy once, and reuse the cached value after.

Yes this is caching the result. Based on profiling we found that computing this every time incurs unacceptable overhead.

tensorrt_llm/_torch/pyexecutor/sampler.py

litaotju · 2025-12-31T07:03:00Z

tensorrt_llm/_torch/pyexecutor/sampler.py

    def _return_log_probs(self, requests: list[LlmRequest]) -> bool:
        return any(req.py_return_log_probs for req in requests)

+    def _prepare_log_probs(self, requests: list[LlmRequest]) -> None:


A style NIT. For function which does not return, I think normally we don't declare -> None, right?

If we write nothing, then it's up to IDE to infer the actual return type. Documentation shows examples that explicitly types -> None: https://docs.python.org/3/library/typing.html#type-aliases

tensorrt_llm/_torch/pyexecutor/sampling_utils.py

litaotju · 2025-12-31T07:12:32Z

tensorrt_llm/_torch/pyexecutor/sampling_utils.py

        return_probs: bool,
        group_metadata: StrategyMetadata | None = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor], float | None]:


When shall we use Optional[float] and when to use float | None?

I think it's like "west const" vs "east const" - it's a convention issue...

As explained before I think we are moving to the new grammar but don't want to suddenly change everything.

tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py

hchings

@tongyuantongyu Could you share some benchmark or Nsight screenshots of the _process_logprobs in the MR description as well? Thanks.

hchings · 2026-01-01T03:36:27Z

tensorrt_llm/_torch/pyexecutor/sampler.py

-    return resolve_sampling_strategy(params, vocab_size=vocab_size)
+    if not hasattr(request, "py_sampling_strategy") or _get_max_beam_width(request) > 1:
+        params = _request_get_sampling_params(request)
+        request.py_sampling_strategy = resolve_sampling_strategy(params, vocab_size=vocab_size)


I think this is to avoid repeated computation by only compute and store on request.py_sampling_strategy once, and reuse the cached value after.

tensorrt_llm/_torch/pyexecutor/sampling_utils.py

zhenhuaw-me · 2026-01-02T01:57:21Z

tensorrt_llm/sampling_params.py

    n: int = 1
    best_of: Optional[int] = None
    use_beam_search: bool = False
+    logprobs_mode: LogprobMode = LogprobMode.RAW


Suggested change

logprobs_mode: LogprobMode = LogprobMode.RAW

Move close to other logprob parameters.

zhenhuaw-me · 2026-01-02T02:01:45Z

tensorrt_llm/sampling_params.py


    # Keep the below fields in sync with tllme.OutputConfig
    logprobs: Optional[int] = None
    prompt_logprobs: Optional[int] = None


Suggested change

prompt_logprobs: Optional[int] = None

logprobs_mode: LogprobMode = LogprobMode.RAW

# Keep the below fields in sync with tllme.OutputConfig

We don't want to add it to the legacy binding type tllme.OutputConfig.

.pre-commit-config.yaml

pyproject.toml

tensorrt_llm/_torch/pyexecutor/llm_request.py

tensorrt_llm/_torch/pyexecutor/sampler.py

- Iintroduces a new optional parameter, logprobs_mode, to the SamplingParams and LlmRequest classes, allowing users to specify the mode of log probabilities to return. - Create process_logprobs function to remove logprobs processing code from process_requests. - add batching based on logprobs_mode to sample_batched_by_strategy - additionally return processed logits from sampling Signed-off-by: Stefan Niebler <[email protected]>

fix step remove dependency on return_genereation_logits align API across backends, add tests test fix group sampling strategy Signed-off-by: Erin Ho <[email protected]> Signed-off-by: Yuan Tong <[email protected]>

Signed-off-by: Erin Ho <[email protected]> Signed-off-by: Yuan Tong <[email protected]>

- Expand test_logits_logprobs to perform a check for processed logprobs - Fix processed logprobs for greedy sampling and when using temperature Signed-off-by: Stefan Niebler <[email protected]>

…chSampler - Updated tensor allocation in TorchSampler to use pinned memory for improved performance during D2H copies. - Modified test_sampled_token_always_in_logprobs to include logprobs_mode parameter for enhanced testing of log probabilities. Signed-off-by: Stefan Niebler <[email protected]>

…lation when not needed - Added LogProbsMode class to define modes for log probabilities: RAW and PROCESSED. - Updated SamplingParams and LlmRequest to utilize LogProbsMode for logprobs_mode parameter. - Enhanced validation to check logprobs_mode against LogProbsMode values. - Modified TorchSampler and related classes to support new logprobs_mode functionality. - Modified TorchSampler to only calculate logprobs when a request needs it - Updated tests to cover new logprobs_mode behavior and ensure correct processing of log probabilities. Signed-off-by: Stefan Niebler <[email protected]>

… logprobs handling - Added max_topk_logprobs parameter to AutoDeployConfig and LlmRequest to control the number of top-k logprobs storable for each token. - Updated TorchSampler to accommodate max_topk_logprobs in logprobs processing and storage. - Enhanced logprobs handling in the sampling process to support both sampled and top-k logprobs. - Enabled batched processing of logprobs to enhance logprobs performance - Modified tests to validate the new max_topk_logprobs functionality and ensure correct logprobs output. Signed-off-by: Stefan Niebler <[email protected]> Signed-off-by: Yuan Tong <[email protected]>

Signed-off-by: Yuan Tong <[email protected]>

…gy when using beam search Signed-off-by: Stefan Niebler <[email protected]>

…t indices instead of a boolean tensor Signed-off-by: Stefan Niebler <[email protected]>

…ch.testing.assert_close for improved readability and precision Signed-off-by: Stefan Niebler <[email protected]>

Signed-off-by: Yuan Tong <[email protected]>

tongyuantongyu · 2026-01-02T07:24:40Z

@litaotju @hchings

This is the Nsight Systems report of a 7B model (Qwen2-7B) running on H100, using Ray (MPI is usually better), processing ~500 requests. I think this is good enough looking at GPU utilization.

hchings mentioned this pull request Dec 5, 2025

[None][fix] Change PyT to always include sampled logprob #9374

Closed

1 task

stnie force-pushed the develop/sampler/refactor_logprobs branch from 0b21e18 to d64119a Compare December 5, 2025 16:27

hchings reviewed Dec 5, 2025

View reviewed changes

tensorrt_llm/sampling_params.py Outdated Show resolved Hide resolved

hchings reviewed Dec 5, 2025

View reviewed changes

tests/unittest/_torch/sampler/test_logits_logprobs.py Outdated Show resolved Hide resolved

hchings mentioned this pull request Dec 5, 2025

[TRTLLM-9736][feat] AsyncLLM and verl integ #9353

Merged

1 task

Superjomn changed the title ~~[None][feat] Add processed logprobs functionality to TorchSampler~~ [TRTLLM-9735][feat] Add processed logprobs functionality to TorchSampler Dec 7, 2025

stnie force-pushed the develop/sampler/refactor_logprobs branch from 4aa1d0c to bfb2195 Compare December 8, 2025 11:04

Funatiq requested review from Funatiq and ixlmar December 8, 2025 12:42

ixlmar reviewed Dec 8, 2025

View reviewed changes

stnie force-pushed the develop/sampler/refactor_logprobs branch from bfb2195 to 18531f1 Compare December 9, 2025 15:50

Funatiq reviewed Dec 10, 2025

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

joyang-nv requested a review from tongyuantongyu December 11, 2025 07:22

tongyuantongyu reviewed Dec 11, 2025

View reviewed changes

tensorrt_llm/_torch/pyexecutor/sampler.py Outdated Show resolved Hide resolved

tensorrt_llm/_torch/pyexecutor/sampler.py Outdated Show resolved Hide resolved

stnie force-pushed the develop/sampler/refactor_logprobs branch 3 times, most recently from d66c983 to 5f75078 Compare December 12, 2025 16:59

tongyuantongyu force-pushed the develop/sampler/refactor_logprobs branch from 64e640c to 3ff5537 Compare December 30, 2025 10:57

litaotju reviewed Dec 31, 2025

View reviewed changes

tensorrt_llm/sampling_params.py Show resolved Hide resolved

litaotju reviewed Dec 31, 2025

View reviewed changes

tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py Show resolved Hide resolved

litaotju reviewed Dec 31, 2025

View reviewed changes

hchings approved these changes Jan 2, 2026

View reviewed changes

zhenhuaw-me added the Decoding/Sampling <NV>Token sampling algorithms in TRTLLM for text gen (top-k, top-p, beam). label Jan 2, 2026

zhenhuaw-me reviewed Jan 2, 2026

View reviewed changes

stnie and others added 17 commits January 2, 2026 14:31

sampled logprob

78ab6f2

fix step remove dependency on return_genereation_logits align API across backends, add tests test fix group sampling strategy Signed-off-by: Erin Ho <[email protected]> Signed-off-by: Yuan Tong <[email protected]>

update

e779211

Signed-off-by: Erin Ho <[email protected]> Signed-off-by: Yuan Tong <[email protected]>

[TRTLLM-9686][feat] Fix issues with processed logprobs functionality.

22c240f

- Expand test_logits_logprobs to perform a check for processed logprobs - Fix processed logprobs for greedy sampling and when using temperature Signed-off-by: Stefan Niebler <[email protected]>

cache sampling_strategy

b4ae3a0

Signed-off-by: Yuan Tong <[email protected]>

fix apply_embedding_bias perf

fc3c89d

Signed-off-by: Yuan Tong <[email protected]>

fix sample_batched_by_strategy perf

2051ffb

Signed-off-by: Yuan Tong <[email protected]>

[TRTLLM-9686][chore] Added skip to testcase and fixed _request_strate…

8dd4a7a

…gy when using beam search Signed-off-by: Stefan Niebler <[email protected]>

[TRTLLM-9686][chore] Refactored speculation_needs_probs to use reques…

12d515d

…t indices instead of a boolean tensor Signed-off-by: Stefan Niebler <[email protected]>

Refactor logprob assertions in test_processed_logprobs_e2e to use tor…

7847268

…ch.testing.assert_close for improved readability and precision Signed-off-by: Stefan Niebler <[email protected]>

Nits, Update to ruff for test file

a6a1356

Signed-off-by: Yuan Tong <[email protected]>

Update logprobs_mode definition

348a547

Signed-off-by: Yuan Tong <[email protected]>

Avoid unnecessary initialization for runtime buffers

c7d6f4f

Signed-off-by: Yuan Tong <[email protected]>

Address reviews

d416a39

Signed-off-by: Yuan Tong <[email protected]>

tongyuantongyu force-pushed the develop/sampler/refactor_logprobs branch from 3ff5537 to d416a39 Compare January 2, 2026 06:54


	prompt_logprobs: Optional[int] = None
	logprobs_mode: LogprobMode = LogprobMode.RAW

[TRTLLM-9735][feat] Add processed logprobs functionality to TorchSampler #9675

Are you sure you want to change the base?

[TRTLLM-9735][feat] Add processed logprobs functionality to TorchSampler #9675

Uh oh!

Conversation

stnie commented Dec 3, 2025 • 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

kill

skip

reuse-pipeline

Uh oh!

stnie commented Dec 3, 2025

Uh oh!

tensorrt-cicd commented Dec 3, 2025

Uh oh!

tensorrt-cicd commented Dec 3, 2025

Uh oh!

stnie commented Dec 5, 2025

Uh oh!

tensorrt-cicd commented Dec 5, 2025

Uh oh!

tensorrt-cicd commented Dec 5, 2025

Uh oh!

Uh oh!

Uh oh!

stnie commented Dec 8, 2025

Uh oh!

tensorrt-cicd commented Dec 8, 2025

Uh oh!

ixlmar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Dec 8, 2025

Uh oh!

stnie commented Dec 9, 2025

Uh oh!

tensorrt-cicd commented Dec 9, 2025

Uh oh!

tensorrt-cicd commented Dec 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

stnie commented Dec 12, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

litaotju commented Dec 31, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

stnie commented Dec 3, 2025 •

edited by coderabbitai bot

Loading

tongyuantongyu commented Jan 2, 2026 •

edited

Loading