The refactor adds a reusable fixture for theLLM instance, reducing test time by about 3x.
The resulting changes in execution revealedhttps://nvbugs/5573238 andhttps://nvbugs/5577178, this PR includes workarounds for both.
Working aroundhttps://nvbugs/5577178 without increasing test time again makes use ofcache_salt, which was supported forLLM.generate_async but not forLLM.generate. The PR resolves this discrepancy.
Also,test_return_logits.py appears not to have been included in the CI tests.

Test Coverage

Change pertains to test.cache_salt is used by existing tests.

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 FollowsTRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (seetest instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
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.

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: DoesNOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: DoesNOT 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: DoesNOT 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: DoesNOT update GitHub pipeline status.

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

--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: DoesNOT 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 thestage-list parameter to access the appropriate container environment. Note: DoesNOT update GitHub check status.

For guidance on mapping tests to stage names, seedocs/source/reference/ci-overview.md
and thescripts/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.

Summary by CodeRabbit

New Features
- Added an optional cache-salting control to generation requests, allowing finer control over KV cache reuse across batched requests. Behavior is unchanged when not used.
Tests
- Expanded coverage for return-logits (sync and async) with scenarios spanning cache reuse and logprob validation.
- Added new integration test entry to the PyTorch suite.
- Improved test stability with targeted timeouts and known-issue handling.

ixlmar force-pushed thetest/refactor-return-logits branch froma7c730f toed6de31Compare

October 13, 2025 09:28

ixlmar changed the title~~[TRTLLM-8551] feat: add cache_salt in LLM.generate and refactor test_return_logits.py~~[TRTLLM-8551][feat] add cache_salt in LLM.generate and refactor test_return_logits.py

Oct 13, 2025

ixlmar force-pushed thetest/refactor-return-logits branch fromed6de31 to7f0996dCompare

October 13, 2025 09:33

Copy link

CollaboratorAuthor

ixlmar commentedOct 13, 2025

/bot run --stage-list "A30-PyTorch-1,A30-PyTorch-2"

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21191 [ run ] triggered by Bot

ixlmar force-pushed thetest/refactor-return-logits branch from7f0996d tof8bfeeeCompare

October 13, 2025 10:55

Copy link

CollaboratorAuthor

ixlmar commentedOct 13, 2025

/bot run --stage-list "A30-PyTorch-1"

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21191 [ run ] completed with stateSUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #15997(Partly Tested) completed with status: 'FAILURE'

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21206 [ run ] triggered by Bot

ixlmar force-pushed thetest/refactor-return-logits branch fromf8bfeee toa81e6e1Compare

October 13, 2025 11:11

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21206 [ run ] completed with stateSUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #16009(Partly Tested) completed with status: 'SUCCESS'

Copy link

CollaboratorAuthor

ixlmar commentedOct 13, 2025

/bot run

ixlmar marked this pull request as ready for review

October 13, 2025 12:34

ixlmar requested a review froma team as acode owner

October 13, 2025 12:34

ixlmar requested review fromFunatiq andhchings

October 13, 2025 12:34

Copy link

Contributor

coderabbitaibot commentedOct 13, 2025

📝 Walkthrough

Walkthrough

Adds an optional cache_salt parameter to BaseLLM.generate and threads it through the sync/async generation path to the executor as cache_salt_id. Updates integration test list. Introduces/expands unit tests for return logits and cache behavior, including parametrization, fixtures, and handling known timeout issues.

Changes

Cohort / File(s)	Summary
LLM API: cache salt threading `tensorrt_llm/llmapi/llm.py`	Adds optional parameter`cache_salt` to`BaseLLM.generate(...)` and passes it through the request pipeline to the executor as`cache_salt_id`. Docstring updated accordingly.
Integration tests list `tests/integration/test_lists/test-db/l0_a30.yml`	Adds`unittest/_torch/sampler/test_return_logits.py` to the l0_a30 integration test set.
Unit tests: logits, cache reuse, async `tests/unittest/_torch/sampler/test_return_logits.py`	Adds fixtures (e.g., cache salter, logits gatherers), parametrized tests over`reuse_cache` and`return_log_probs`, uses`cache_salt` in sync/async`llm.generate`, validates context/generation logits and logprobs, includes xfail handling for known timeout bugs, sets`KvCacheConfig(..., enable_block_reuse=True)`.

Sequence Diagram(s)

sequenceDiagram  autonumber  participant C as Client  participant L as BaseLLM  participant A as Async Gen Pipeline  participant E as Executor  Note over C,L: Synchronous call with optional cache_salt  C->>L: generate(inputs, ..., cache_salt)  L->>A: submit(inputs, ..., cache_salt)  A->>E: execute(request, cache_salt_id=cache_salt)  E-->>A: tokens, logits, caches  A-->>L: RequestOutput  L-->>C: RequestOutput

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 8.33% which is insufficient. The required threshold is 80.00%.	You can run`@coderabbitai generate docstrings` to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The title follows the repository conventions by including the JIRA ticket identifier, the change type, and a concise description of both the addition of the cache_salt parameter and the related test refactor in a single clear phrase.
Description Check	✅ Passed	The pull request description includes the required ## Description, ## Test Coverage, and ## PR Checklist sections from the repository template, clearly explains the issue and solution, and provides guidance for CI interaction.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

coderabbitaibot reviewed

Oct 13, 2025

View reviewed changes

Copy link

Contributor

coderabbitaibot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 0

🧹 Nitpick comments (2)

tests/unittest/_torch/sampler/test_return_logits.py (2)
38-58:Consider using a more robust salt generation approach.
TheCacheSalter class uses a class variable_salt that is mutated across test invocations. While this works for serial test execution, it could lead to non-deterministic behavior if tests are run in parallel or the module is imported multiple times.
Consider using a more explicit approach, such as a fixture-scoped counter or UUID generation:
 class CacheSalter:--    _salt = 0-     @classmethod-    def get_salt_unique(cls) -> str:-        cls._salt += 1-        return str(cls._salt)+    def get_salt_unique(cls, iteration: int) -> str:+        return str(iteration)
Or use UUIDs for guaranteed uniqueness:
importuuidclassCacheSalter:@classmethoddefget_salt_unique(cls)->str:returnstr(uuid.uuid4())
However, the current implementation is acceptable given that pytest runs module tests serially by default.
138-156:Add explanatory comment for xfail condition
Above thepytest.xfail("Known bug: https://nvbugs/5577178") at line 152, include:
# Bug manifests when context_logits gathering is enabled with cache reuse
to clarify why onlygather_context_logits and reuse_cache triggers xfail.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and betweend145e87 anda81e6e1.

📒 Files selected for processing (3)

tensorrt_llm/llmapi/llm.py (3 hunks)
tests/integration/test_lists/test-db/l0_a30.yml (1 hunks)
tests/unittest/_torch/sampler/test_return_logits.py (2 hunks)

🧰 Additional context used

📓 Path-based instructions (3)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use only spaces, no tabs; indent with 4 spaces.

Files:

tensorrt_llm/llmapi/llm.py
tests/unittest/_torch/sampler/test_return_logits.py

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.

Files:

tensorrt_llm/llmapi/llm.py
tests/unittest/_torch/sampler/test_return_logits.py

**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).

Files:

tensorrt_llm/llmapi/llm.py
tests/unittest/_torch/sampler/test_return_logits.py

🧬 Code graph analysis (1)

tests/unittest/_torch/sampler/test_return_logits.py (5)

tests/unittest/_torch/sampler/test_beam_search.py (1)
llm (37-49)
tests/unittest/_torch/sampler/test_best_of_n.py (1)
llm (36-42)
tensorrt_llm/llmapi/llm.py (3)
LLM (1087-1103)
generate (249-330)
generate_async (333-498)
tensorrt_llm/scaffolding/task.py (1)
logprobs (99-100)
tensorrt_llm/_torch/pyexecutor/llm_request.py (1)
generation_logits (298-307)

🔇 Additional comments (7)

tensorrt_llm/llmapi/llm.py (1)
265-265:LGTM! cache_salt parameter properly added to synchronous generate.
Thecache_salt parameter is correctly added to the synchronousgenerate method signature, documented, and threaded through togenerate_async (which already supported it). This enables KV cache salting for batched synchronous generation.
Also applies to: 286-286, 317-318
tests/integration/test_lists/test-db/l0_a30.yml (1)
24-24:LGTM! Test file added to integration suite.
The new test file is appropriately added to the PyTorch test block in the A30 integration test list.
tests/unittest/_torch/sampler/test_return_logits.py (5)
12-15:LGTM! KV cache configuration updated for block reuse.
The global KV cache config now enables block reuse, which is necessary to test the cache_salt functionality introduced in this PR.
18-36:LGTM! Fixture parametrization enables comprehensive test coverage.
The module-scoped fixtures parametrize key configuration options (gather_context_logits, gather_generation_logits, disable_overlap_scheduler, sampler_type). This creates a test matrix that exercises different combinations, which is appropriate for thorough testing of the return_logits feature.
87-107:Workaround for LLM shutdown timeout is acceptable but should be tracked.
The monkeypatch that converts pytest-timeout failures to xfails is a pragmatic workaround for the known bug (nvbugs/5577178). However, this approach could mask real test failures if other timeout issues arise.
Consider adding a more specific check to ensure only the expected shutdown timeout is caught:
def_exit_with_xfail_on_timeout(self,exc_type,exc_value,traceback)->bool:import_pytest.outcomestry:returnold_exit(self,exc_type,exc_value,traceback)except_pytest.outcomes.Failedase:ife.msgand"pytest-timeout"ine.msg.lower():# Add more specific check for shutdown contextif"shutdown"instr(e)or"_shutdown"instr(e):pytest.xfail("Known LLM shutdown issue (https://nvbugs/5577178).")raise# Re-raise if it's a different timeoutelse:raise
This would help ensure that only shutdown-related timeouts are converted to xfails, while other timeout failures are still caught.
110-116:Test markers appropriately configured for known issues.
The timeout and threadleak markers are correctly added to handle the known LLM shutdown issue (nvbugs/5577178). The 120-second timeout is reasonable, and the signal method ensures child processes are also terminated.
200-234:Unconditional xfail for async generation_logits is appropriate.
The async test streams one token at a time, so the nvbugs/5573238 mismatch affects all streaming cases and warrants the unconditional xfail.

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21214 [ run ] triggered by Bot

Funatiq approved these changes

Oct 13, 2025

View reviewed changes

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21214 [ run ] completed with stateSUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #16015 completed with status: 'FAILURE'

ixlmar force-pushed thetest/refactor-return-logits branch froma81e6e1 to90b1100Compare

October 13, 2025 15:42

Copy link

CollaboratorAuthor

ixlmar commentedOct 13, 2025

/bot run --disable-fail-fast

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21235 [ run ] triggered by Bot

Copy link

Collaborator

tensorrt-cicd commentedOct 13, 2025

PR_Github #21235 [ run ] completed with stateSUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #16028 completed with status: 'FAILURE'

ixlmar force-pushed thetest/refactor-return-logits branch 2 times, most recently from0f7d268 to518f2f6Compare

October 14, 2025 07:32

ixlmar requested a review froma team as acode owner

October 14, 2025 07:32

ixlmar requested a review fromnetanel-haber

October 14, 2025 07:32

Copy link

CollaboratorAuthor

ixlmar commentedOct 14, 2025

/bot run --disable-fail-fast

Copy link

Collaborator

tensorrt-cicd commentedOct 14, 2025

PR_Github #21334 [ run ] completed with stateFAILURE

Copy link

CollaboratorAuthor

ixlmar commentedOct 14, 2025

/bot run --disable-fail-fast

Copy link

Collaborator

tensorrt-cicd commentedOct 14, 2025

PR_Github #21335 [ run ] triggered by Bot

syuoni reviewed

Oct 14, 2025

View reviewed changes

tensorrt_llm/llmapi/llm.py OutdatedShow resolvedHide resolved

ixlmar force-pushed thetest/refactor-return-logits branch 2 times, most recently from10fbddd toe1a2d86Compare

October 14, 2025 14:47

DomBrown approved these changes

Oct 14, 2025

View reviewed changes

Copy link

CollaboratorAuthor

ixlmar commentedOct 14, 2025

/bot run --disable-fail-fast

Copy link

Collaborator

tensorrt-cicd commentedOct 14, 2025

PR_Github #21367 [ run ] triggered by Bot

Copy link

Collaborator

tensorrt-cicd commentedOct 14, 2025

PR_Github #21335 [ run ] completed with stateABORTED
LLM/main/L0_MergeRequest_PR #16104 (Blue Ocean) completed with status: ABORTED

ixlmar requested a review fromsyuoni

October 14, 2025 15:03

syuoni approved these changes

Oct 14, 2025

View reviewed changes

Copy link

Collaborator

syuoni left a comment

Choose a reason for hiding this comment

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

LGTM, please update the doc string and API reference, thanks!

tensorrt_llm/llmapi/llm.py OutdatedShow resolvedHide resolved

feat: add cache_salt in LLM.generate and refactor test_return_logits.py

54893d0

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

ixlmar force-pushed thetest/refactor-return-logits branch frome1a2d86 to54893d0Compare

October 14, 2025 15:53

Copy link

CollaboratorAuthor

ixlmar commentedOct 14, 2025

/bot run --disable-fail-fast

Copy link

Collaborator

tensorrt-cicd commentedOct 14, 2025

PR_Github #21373 [ run ] triggered by Bot

Copy link

Collaborator

tensorrt-cicd commentedOct 14, 2025

PR_Github #21367 [ run ] completed with stateABORTED
LLM/main/L0_MergeRequest_PR #16134 (Blue Ocean) completed with status: ABORTED

Copy link

Collaborator

tensorrt-cicd commentedOct 15, 2025

PR_Github #21373 [ run ] completed with stateSUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #16139 completed with status: 'FAILURE'

Copy link

CollaboratorAuthor

ixlmar commentedOct 15, 2025

/bot run

Copy link

Collaborator

tensorrt-cicd commentedOct 15, 2025

PR_Github #21444 [ run ] triggered by Bot

dcamporaenabled auto-merge (squash)

October 15, 2025 09:52

Copy link

Collaborator

tensorrt-cicd commentedOct 15, 2025

PR_Github #21444 [ run ] completed with stateSUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #16192 completed with status: 'SUCCESS'

dcampora merged commit0510b34 intoNVIDIA:main

Oct 15, 2025

5 checks passed

ixlmar deleted the test/refactor-return-logits branch

October 15, 2025 09:56

govind-ramnarayan pushed a commit to nv-auto-deploy/TensorRT-LLM that referenced this pull request

Oct 21, 2025