[https://nvbugs/5708901][perf] avoid logits copies when computing logprobs - #16958
Conversation
26337c0 to
93c4651
Compare
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
93c4651 to
bf4109d
Compare
|
/bot run |
|
PR_Github #62231 [ run ] triggered by Bot. Commit: |
|
PR_Github #62231 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62454 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe sampler separates raw and processed logprob requirements, excludes beam-search requests from shared buffers, uses output-buffered log-softmax fusion, and adds comprehensive mixed-batch validation. ChangesSampler logprob refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Grouper as _CachingRequestGrouper
participant Sampler as _sample_batched_by_strategy
participant Processor as _process_logprobs
participant Fusion as Fusions
Grouper->>Sampler: grouped requests and raw-logprob mask
Sampler->>Sampler: filter beam requests and compute logprob buffers
Sampler->>Processor: batched result with raw and processed indices
Processor->>Fusion: gather_log_softmax_with_output
Fusion->>Processor: write raw logprobs into output buffer
Processor->>Processor: gather and scatter final logprob outputs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/sampler/test_torch_sampler.py (1)
2510-2518: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winTest coverage summary (tests/ review).**
- Changed test functions: only
test_unbatch_sampling_resultswas touched, and solely to follow the_BatchedSamplingResultfield renames (req_indices,next_tokens_cuda_int) plus one added side-effect assertion. No test functions were added or removed.- Test-list files: no entries under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/were modified; since no new test functions were introduced, no list updates are strictly required — but if the logprob paths get dedicated cases, please register them in the appropriatetest-db/list.- Coverage verdict: insufficient for this PR's behavior change. The rewritten logprob pipeline (separate raw vs processed buffers, per-group partial PROCESSED subsets, joint raw+processed indexing,
gather_log_softmax_with_output) has no direct assertions here; the only touched test exercises unbatching. Recommended additions: aTestBatchedSampling-style case parameterized overlogprobs_mode(RAW / PROCESSED / mixed within one strategy group) validating sampled logprob values, ranks and top-k output, including permuted request order and multi-step (draft) rows.I'm happy to draft those cases if useful.
As per path instructions, "Always produce a test coverage summary, even if no issues are found" and the summary must state which tests changed, whether they are listed under
tests/integration/test_lists/, and a coverage verdict.🤖 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/_torch/sampler/test_torch_sampler.py` around lines 2510 - 2518, Extend the tests around TestBatchedSampling to cover the rewritten logprob pipeline, rather than only updating test_unbatch_sampling_results for renamed fields. Parameterize cases for RAW, PROCESSED, and mixed logprobs_mode within one strategy group, asserting sampled logprob values, ranks, and top-k output with permuted request order and multi-step draft rows. Register any new integration test cases in the appropriate tests/integration/test_lists/ test-db or qa list.Source: Path instructions
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
5072-5072: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
logprobs_inout_indices_cuda_sizeis assigned but unused in the 1-beam branch.Line 5072 computes it and the 1-beam block never reads it (only the n-beam block at Line 5147 does). Drop the 1-beam assignment.
Also applies to: 5147-5147
🤖 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/_torch/pyexecutor/sampler/sampler.py` at line 5072, Remove the unused logprobs_inout_indices_cuda_size assignment from the 1-beam branch, while retaining the assignment in the n-beam branch where it is consumed.
🤖 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/_torch/pyexecutor/sampler/sampler.py`:
- Around line 4576-4619: Fix the mixed processed-logprobs branch around
proc_lp_step_mask_cuda to gather rows by the selected request indices rather
than constructing a mask from only proc_lp_steps. Apply the same group-aligned
row selection to logits, softmax, and tensor temperature data, preserving scalar
temperatures, and add coverage for mixed processed/non-processed groups with
permuted request ordering.
- Around line 4995-5013: The 1-beam logprob path must preserve alignment with
rows in logprobs_cuda, whose order follows logprobs_reqs_indices. Update the
filtering logic around logprobs_reqs_indices_1_beam and the subsequent 1-beam
torch.gather flow to carry explicit source row indices or reorder/slice
logprobs_cuda before gathering, ensuring mixed 1-beam and n-beam batches read
the correct request rows.
---
Outside diff comments:
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 2510-2518: Extend the tests around TestBatchedSampling to cover
the rewritten logprob pipeline, rather than only updating
test_unbatch_sampling_results for renamed fields. Parameterize cases for RAW,
PROCESSED, and mixed logprobs_mode within one strategy group, asserting sampled
logprob values, ranks, and top-k output with permuted request order and
multi-step draft rows. Register any new integration test cases in the
appropriate tests/integration/test_lists/ test-db or qa list.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Line 5072: Remove the unused logprobs_inout_indices_cuda_size assignment from
the 1-beam branch, while retaining the assignment in the n-beam branch where it
is consumed.
🪄 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: a4cff641-f0dd-4c40-8715-a2d53e6ed8b7
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_torch_sampler.py
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
PR_Github #62454 [ run ] completed with state
|
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
Note: New tests take ~26 seconds on RTX 6000 PRO. |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
4735-4764: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep RAW beam requests in
raw_logprobs_reqs_indices
Filtering this list down topy_beam_width == 1drops beam requests from_process_logprobs, sosampled_log_prob_indicesis never refreshed for RAW-mode beam batches and can be copied out stale. Keep the full request list here and use the filtered tensor only forbatch_raw_logprob_indices_cuda.🤖 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/_torch/pyexecutor/sampler/sampler.py` around lines 4735 - 4764, In the raw-logprob preparation block, preserve all entries in reqs_indices_needing_raw_logprobs for _process_logprobs and sampled_log_prob_indices updates, including beam requests. Keep filtering py_beam_width == 1 only when building gather_raw_logprobs_req_indices_tensor and batch_raw_logprob_indices_cuda, while assigning the unfiltered request indices to reqs_indices_needing_raw_logprobs.
🧹 Nitpick comments (1)
tests/unittest/_torch/sampler/test_logits_logprobs.py (1)
916-944: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe two
FIXME/exclusion comments carve out exactly the paths changed in this PR.
- Line 932-934 skips
temperature+logprobs=2whenever the sampler is beam-capable, becausevalidate_requestrejectsnum_logprobs > 1for any request oncemax_beam_width > 1— even single-beam ones.- Line 942-944 skips PROCESSED + beam-capable batches, which is precisely the branch containing the beam-filter defect flagged in
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py(Lines 4568-4602).Want me to open a tracking issue for the
validate_requestover-restriction, and/or add a mixed PROCESSED + beam-search case once the sampler branch is fixed?🤖 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/_torch/sampler/test_logits_logprobs.py` around lines 916 - 944, Update LOGPROBS_CASES and BASE_CASES so the tests no longer exclude temperature beam-capable requests with logprobs=2 or beam-search requests using processed logprobs. Enable coverage for these paths, including the mixed PROCESSED plus beam-search case, and remove the obsolete FIXME-based filters.
🤖 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/_torch/pyexecutor/sampler/sampler.py`:
- Around line 4568-4602: Correct the beam-search filtering in the
`group_gather_processed_logprobs` setup: retain `py_beam_width == 1` requests
for gathering and exclude the beam-search requests, rather than clearing the
width-one rows. Use each request’s actual group index (`req_idx`) when updating
the mask, not its enumeration position in
`need_processed_logprobs_req_indices_list`, and keep the count aligned with the
resulting gathered rows.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 4735-4764: In the raw-logprob preparation block, preserve all
entries in reqs_indices_needing_raw_logprobs for _process_logprobs and
sampled_log_prob_indices updates, including beam requests. Keep filtering
py_beam_width == 1 only when building gather_raw_logprobs_req_indices_tensor and
batch_raw_logprob_indices_cuda, while assigning the unfiltered request indices
to reqs_indices_needing_raw_logprobs.
---
Nitpick comments:
In `@tests/unittest/_torch/sampler/test_logits_logprobs.py`:
- Around line 916-944: Update LOGPROBS_CASES and BASE_CASES so the tests no
longer exclude temperature beam-capable requests with logprobs=2 or beam-search
requests using processed logprobs. Enable coverage for these paths, including
the mixed PROCESSED plus beam-search case, and remove the obsolete FIXME-based
filters.
🪄 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: 0d5b8126-cfc1-4753-b43f-e3c72ca91631
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_logits_logprobs.py
…opies Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
3126-3154: 🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy liftProcessed-logprobs beam mask is inverted and indexed in the wrong space.
Two problems in the
self._use_beam_searchbranch:
- Polarity:
gather_processed_logprobs_group_req_indices_listcollects thepy_beam_width == 1requests — the ones that should contribute rows — and then clears them in the mask, leaving the beam-search requests selected.- Index space:
grp_idxis a position withinneed_processed_logprobs_req_indices_list(positions among the mask'sTrueentries), not a position within the group-length maskgroup_gather_processed_logprobs. It only coincides when the selected requests happen to form a prefix of the group.Consequently
num_gather_processed_logprobs_req_indices(count of 1-beam requests) disagrees with the rows the mask selects, so mixed beam/non-beam groups can route the wrong logit rows intobatch_processed_logprobs_cuda.🐛 Suggested fix: build the mask from group positions and keep 1-beam requests
- if self._use_beam_search: - gather_processed_logprobs_group_req_indices_list = [ - grp_idx - for grp_idx, req_idx in enumerate( - need_processed_logprobs_req_indices_list - ) - if requests[req_idx].py_beam_width == 1 - ] - - # repurpose group_need_processed_logprobs - group_gather_processed_logprobs = group_need_processed_logprobs - del group_need_processed_logprobs - gather_processed_logprobs_group_req_indices = torch.tensor( - gather_processed_logprobs_group_req_indices_list, - dtype=need_processed_logprobs_req_indices.dtype, - ) - group_gather_processed_logprobs[ - gather_processed_logprobs_group_req_indices - ] = False - num_gather_processed_logprobs_req_indices = len( - gather_processed_logprobs_group_req_indices_list - ) + if self._use_beam_search: + # Positions within the group (not within the selected sublist) whose + # request is beam-search and must therefore be excluded. + drop_group_positions = [ + grp_pos + for grp_pos, req_idx in enumerate(group_req_indices.tolist()) + if requests[req_idx].py_beam_width > 1 + ] + group_gather_processed_logprobs = group_need_processed_logprobs.clone() + del group_need_processed_logprobs + if drop_group_positions: + group_gather_processed_logprobs[ + torch.tensor(drop_group_positions, dtype=torch.long) + ] = False + num_gather_processed_logprobs_req_indices = cast( + int, group_gather_processed_logprobs.sum().item() + )Please also cover a group mixing beam-search and single-beam
PROCESSED-logprob requests with permuted ordering.🤖 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/_torch/pyexecutor/sampler/sampler.py` around lines 3126 - 3154, Fix the self._use_beam_search branch so group_gather_processed_logprobs selects only py_beam_width == 1 requests. Build the mask using each request’s position in the full group mask, not grp_idx within need_processed_logprobs_req_indices_list; preserve those group positions as true and clear the beam-search positions. Keep num_gather_processed_logprobs_req_indices equal to the number of retained single-beam requests, and cover mixed PROCESSED-logprob requests in permuted order.
🧹 Nitpick comments (2)
tests/unittest/_torch/sampler/test_logits_logprobs.py (1)
899-899: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNit: class-name typo and missing return annotations.
TestLogsprobsInBatchedSampling→TestLogprobsInBatchedSampling(it leaks into every generated test ID), and_validate_logprobs/test_logprobs/ the innervalidate_logprob_and_rankclosure lack-> None.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".Also applies to: 1285-1291, 1408-1417
🤖 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/_torch/sampler/test_logits_logprobs.py` at line 899, Rename TestLogsprobsInBatchedSampling to TestLogprobsInBatchedSampling, and add -> None annotations to _validate_logprobs, test_logprobs, and the inner validate_logprob_and_rank closure, including the corresponding functions in the referenced sections.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py (1)
321-340: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the
outcontract. A short Google-style docstring ongather_log_softmax_with_outputstating the expected shape (indices_cuda.size(0) x vocab), thatoutmust not aliasinputs_cuda, and that dtype conversion is implicit would help future callers.🤖 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/_torch/pyexecutor/sampler/ops/vanilla.py` around lines 321 - 340, Add a concise Google-style docstring to Fusions.gather_log_softmax_with_output documenting that out has shape indices_cuda.size(0) × vocab, must not alias inputs_cuda, and receives implicit dtype conversion. Leave the implementation unchanged.Source: Coding guidelines
🤖 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 `@tests/unittest/_torch/sampler/test_logits_logprobs.py`:
- Around line 1140-1149: Update the sparsity assertion in the sequence-slot
setup so it requires num_seq_slots to be at most half of allowed_slots.size,
preserving rng.choice’s replace=False validity while enforcing the intended
sparse assignment.
---
Duplicate comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 3126-3154: Fix the self._use_beam_search branch so
group_gather_processed_logprobs selects only py_beam_width == 1 requests. Build
the mask using each request’s position in the full group mask, not grp_idx
within need_processed_logprobs_req_indices_list; preserve those group positions
as true and clear the beam-search positions. Keep
num_gather_processed_logprobs_req_indices equal to the number of retained
single-beam requests, and cover mixed PROCESSED-logprob requests in permuted
order.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py`:
- Around line 321-340: Add a concise Google-style docstring to
Fusions.gather_log_softmax_with_output documenting that out has shape
indices_cuda.size(0) × vocab, must not alias inputs_cuda, and receives implicit
dtype conversion. Leave the implementation unchanged.
In `@tests/unittest/_torch/sampler/test_logits_logprobs.py`:
- Line 899: Rename TestLogsprobsInBatchedSampling to
TestLogprobsInBatchedSampling, and add -> None annotations to
_validate_logprobs, test_logprobs, and the inner validate_logprob_and_rank
closure, including the corresponding functions in the referenced sections.
🪄 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: 28f6e047-e50e-40bd-8765-bbf91af50a7e
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_logits_logprobs.pytests/unittest/_torch/sampler/test_torch_sampler.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unittest/_torch/sampler/test_torch_sampler.py
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
/bot run |
|
PR_Github #62841 [ run ] triggered by Bot. Commit: |
|
PR_Github #62841 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62967 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with some little comments.
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
PR_Github #62967 [ run ] completed with state
|
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62998 [ run ] triggered by Bot. Commit: |
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #63025 [ run ] triggered by Bot. Commit: |
|
PR_Github #62998 [ run ] completed with state |
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #63039 [ run ] triggered by Bot. Commit: |
|
PR_Github #63025 [ run ] completed with state |
|
PR_Github #63039 [ run ] completed with state |
stnie
left a comment
There was a problem hiding this comment.
Seems fine. Only a few small questions.
Do you have any measurements, how much this improves performance?
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. 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. |
|
/bot run --disable-fail-fast |
|
PR_Github #63057 [ run ] triggered by Bot. Commit: |
|
PR_Github #63057 [ run ] completed with state |
…probs (#16958) Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
Description
Logprobs handling improvements in
TorchSampler:Avoid snapshotting logits for later computing log probs, but either
track logit slice indices (for computing raw logprobs); or
save one
log_softmax()call by directly computing logprobs by offsetting logits by log(max(SoftMax(logits))) (for processed logprobs)Simplify indexing in processed logprobs calculation (replace
_Packed...withtorch.repeat_interleaved+.masked_select)Fast-path indexing in processed logprobs calculation if all requests require processed logprobs
Document several limitations/assumptions and potential future optimizations in code comments
Test Coverage
Relying on 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 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.Dev Engineer Review
TorchSamplerlogprobs computation to reduce work and avoid logits snapshots:logprobs_cudapath)._CachingRequestGrouper.group_requests_by_strategy_keyto returnneed_raw_logprobs, and reshaped_BatchedSamplingResultto carry separate processed/raw indices/logit indices._process_logprobsto acceptlogits_cudaandnew_tokens_cudaexplicitly and adjusted thelogprobs_cuda/store layout assumptions._process_requestscall site and invocation ordering around_unbatch_sampling_results.tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py:Fusions.gather_scatter/_gather_scatter_impl.Fusions.gather_log_softmax(...)with an output-bufferedFusions.gather_log_softmax_with_output(...), backed by_gather_log_softmax_impl(out=...)that writes in-place (preservingtorch._dynamo.mark_dynamicbehavior).QA Engineer Review
tests/unittest/_torch/sampler/test_torch_sampler.py::test_unbatch_sampling_results_BatchedSamplingResultconstruction/side-effect assertions to usereq_indicesandnext_tokens_cuda_int.tests/unittest/_torch/sampler/test_logits_logprobs.py::TestLogsprobsInBatchedSampling::test_logprobsLogprobMode.PROCESSED, including beam-search variants when supported by constraints; verifies sampled logprobs/top-k behavior matches recomputed expectations and checks logits aren’t mutated (with poisoned buffer regions to catch stale-state reliance).tests/integration/test_lists/test-db/l0_a10.yml:unittest/_torch/sampler/test_torch_sampler.pytests/integration/test_lists/test-db/l0_a30.yml:unittest/_torch/sampler/test_logits_logprobs.py