[TRTLLM-13212][refactor] Unify sampler ops and clean up dead code in TorchSampler#16365
Merged
zhaoyangwang-nvidia merged 5 commits intoJul 15, 2026
Merged
Conversation
…shinfer op variants, drop dead trtllm op Follow-ups promised in PR NVIDIA#15542 review: - vanilla: fold top_k_/top_p_/temperature_sampling_batch thin wrappers into top_k_top_p_sampling_batch via default args (top_k=None disables top-k, top_p=1 disables top-p); drop the pass-through greedy() wrapper (greedy re-export now aliases greedy_search_sampling_batch directly). - flashinfer ops: merge the seed/offset and generator variants of each sampling op into a single signature with keyword-only randomness args, matching flashinfer's own unified API (>= 0.6.4 accepts both; explicit seed/offset takes precedence). 6 wrappers -> 4. - Fix seed/offset type hints on sampling_batch_spec_dec_one_model (Optional[int] -> Optional[torch.Tensor]; callers pass device tensors for CUDA graph compatibility). - Remove dead trtllm::compute_probs_from_logits_op chain (no callers anywhere): Python register_fake, thop wrapper + TORCH_LIBRARY registration, and the exclusive kernel chain in dynamicTreeKernels.cu (topKProbStage1/2, invokeTopKTopPMaskingForProbs, applyTopKTopPForProbOp, computeSoftmaxForProbOp, computeProbsFromLogits) plus now-unused torch/ATen includes. - Refresh stale docstring references in sanitize_top_k. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
… public API The sampling facade (sampling_utils) and sampler.py were reaching into underscore-private members of ops.vanilla. Promote the de-facto public surface and remove a duplicated constant: - _GREEDY_TEMPERATURE_THRESHOLD -> GREEDY_TEMPERATURE_THRESHOLD; document the contract with the spec-decoding metadata layer, and derive DISABLE_TEMP_VAL in speculative/interface.py from it (GREEDY_TEMPERATURE_THRESHOLD / 10 == 1e-5, value unchanged) so the sentinel and the threshold can no longer drift apart silently. - _safely_apply_temperature_inplace -> safely_apply_temperature_inplace. - _Fusions -> Fusions (compiled logprobs fusion helpers used by sampler.py); per-method _impl internals stay private. - Drop the duplicated BEAM_SEARCH_PAD_TOKEN definition in sampling_utils; ops.vanilla now holds the single definition and the facade re-exports it. speculative/interface.py imports the threshold via the sampling_utils facade, keeping sampling_utils the only importer of the ops backend modules. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Repo-wide audit found four methods with zero callers (no dynamic dispatch, no test usage): - _handle_finish_reasons: superseded by _handle_first_finish_reasons / direct _handle_finish_reasons_impl calls. - _longest_stop_word_len: logic re-implemented in the live _check_stop_words_length. - _requests_with_stop_words / _request_indices_with_stop_words: unused. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…th penalty path, document flashinfer op contracts - _get_logprobs_from_request: use forward slicing for the history view; [:, :-preallocate_extra_steps, :] silently yields an empty view when preallocate_extra_steps == 0 (the default), skipping history filling. - _apply_min_length_penalty: return None instead of echoing the in-place-mutated logits; defer the num_steps/num_beams host tensor->list conversion until after the early-out; flatten the double-nested guard ifs into early continues, hoisting the per-request offset accumulation above the guards. - _flashinfer_check_nans: restore the explanatory comments lost in the ops-module split (NVIDIA#15542): the constant False deliberately keeps FlashInfer's syncing check_nan path disabled, while the async device-side assert provides NaN protection without stalling the pipeline (see flashinfer issue NVIDIA#1575). - ops/flashinfer: per-op docstrings pointing at the module-level generator-vs-seed/offset contract; note why the 1:1 pipeline-stage wrappers (mask/softmax/renorm) exist (guarded import, PDL decision). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
unify-sampler-ops
branch
from
July 14, 2026 08:34
68daeac to
5cf0c6a
Compare
zhaoyangwang-nvidia
marked this pull request as ready for review
July 14, 2026 08:35
zhaoyangwang-nvidia
requested review from
schetlur-nv,
tfogal,
tongyuantongyu,
yizhang-nv and
ziyixiong-nv
July 14, 2026 08:35
QiJune
approved these changes
Jul 14, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59173 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59173 [ run ] completed with state
|
The unified flashinfer ops now forward seed/offset (alongside generator) to the underlying flashinfer.sampling.* calls, but the test mocks that patch those functions only accepted generator, so the eager batched-sampling tests failed with 'unexpected keyword argument seed'. Add seed/offset (defaulting to None) to the four affected mock signatures. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Collaborator
Author
|
/bot run |
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59359 [ run ] triggered by Bot. Commit: |
Contributor
|
Caution Review failedFailed to post review comments. GitHub was unavailable or timed out while CodeRabbit was posting the review. Please request a new review later if the pull request still needs one. Use ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (10)
💤 Files with no reviewable changes (3)
⏰ Context from checks skipped due to timeout. (1)
🧰 Additional context used📓 Path-based instructions (6)**/*.py📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Files:
**/*.{cpp,cc,cxx,h,hpp,py}📄 CodeRabbit inference engine (AGENTS.md)
Files:
**/*📄 CodeRabbit inference engine (AGENTS.md)
Files:
tensorrt_llm/**/*.{py,cpp,cc,cxx,h,hpp}📄 CodeRabbit inference engine (AGENTS.md)
Files:
tests/unittest/**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
tests/**⚙️ CodeRabbit configuration file
Files:
🧠 Learnings (14)📚 Learning: 2026-02-13T10:15:37.120ZApplied to files:
📚 Learning: 2026-04-24T23:54:27.755ZApplied to files:
📚 Learning: 2026-05-29T08:18:25.109ZApplied to files:
📚 Learning: 2026-05-29T08:18:27.934ZApplied to files:
📚 Learning: 2026-06-25T13:04:55.284ZApplied to files:
📚 Learning: 2026-03-06T11:45:02.068ZApplied to files:
📚 Learning: 2026-03-09T12:34:56.631ZApplied to files:
📚 Learning: 2026-04-16T00:07:18.998ZApplied to files:
📚 Learning: 2026-04-24T03:28:03.123ZApplied to files:
📚 Learning: 2026-05-16T01:43:01.298ZApplied to files:
📚 Learning: 2026-02-25T01:48:05.078ZApplied to files:
📚 Learning: 2026-05-29T08:18:31.547ZApplied to files:
📚 Learning: 2026-02-27T21:32:25.857ZApplied to files:
📚 Learning: 2025-12-12T03:27:08.565ZApplied to files:
📝 WalkthroughWalkthroughChangesUnified sampling interfaces
Dynamic tree probability operation removal
Fake allocate-output contract
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Sampler
participant SamplingUtils
participant FlashInfer
Sampler->>SamplingUtils: select sampling strategy
SamplingUtils->>FlashInfer: pass logits or probabilities with randomness inputs
FlashInfer-->>SamplingUtils: return sampled token
SamplingUtils-->>Sampler: return token and sampling results
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
|
PR_Github #59359 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59408 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59408 [ run ] completed with state |
ixlmar
reviewed
Jul 15, 2026
| return tokens | ||
|
|
||
|
|
||
| # The three ops below wrap the mask -> softmax -> renorm pipeline stages 1:1. |
Collaborator
There was a problem hiding this comment.
For instance, softmax_op will fail if FlashInfer is not available. Importing something that certainly fails if it ever gets invoked may result in late failures (imagine a service being "healthy" and crashing only after a user request selects some unsupported behavior). It might be preferable if that service failed at startup time already.
Collaborator
Author
| # only called on the flashinfer sampler / speculative-worker paths. | ||
| from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import ( | ||
| sampling_from_probs_generator_op as sampling_from_probs_generator_op, | ||
| sampling_from_probs_op as sampling_from_probs_op, |
Collaborator
There was a problem hiding this comment.
Replace from a import b as b with from a import b (multiple occurrences).
|
|
||
| # Re-export the torch greedy op (used by drafting_loops and speculative/interface). | ||
| greedy = _torch_greedy | ||
| greedy = greedy_search_sampling_batch |
Collaborator
There was a problem hiding this comment.
greedy_search_sampling_batch is a better name. Should we perhaps change the dependent modules to use this name instead?
| ``vocab_size`` (== keep all tokens), leaving genuine top_k values | ||
| untouched. | ||
| """ | ||
| return torch.where(top_k > 0, top_k, torch.full_like(top_k, vocab_size)).clamp(max=vocab_size) |
Collaborator
There was a problem hiding this comment.
Consider using .masked_fill_.
zhaoyangwang-nvidia
added a commit
to zhaoyangwang-nvidia/TensorRT-LLM
that referenced
this pull request
Jul 16, 2026
Apply ixlmar's post-merge review suggestions from NVIDIA#16365: - ops/flashinfer.py: without flashinfer installed, calling a sampling op now raises an ImportError with installation guidance instead of a bare NameError. The module intentionally stays importable (sampling_utils imports it unconditionally and non-flashinfer sampler paths must keep working); startup-time enforcement remains in TorchSampler.__init__. Document that contract in the module docstring. - sampling_utils.py: collapse the "from x import y as y" re-export lines into plain grouped imports plus an explicit __all__ (mypy runs strict, so implicit re-export is disabled). - Drop the "greedy" alias of greedy_search_sampling_batch; use the descriptive name directly in drafting_loops.py and speculative/interface.py. - sanitize_top_k: replace torch.where(...).clamp(...) with clamp(...).masked_fill_(...), avoiding the full_like intermediate. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes
Refactor
Description
Follow-ups promised in the #15542 review (sampler ops centralization), plus dead-code
removal and small correctness/readability fixes found while auditing the sampler module.
All changes are behavior-preserving unless noted.
Dedup thin wrappers (ops/vanilla.py) — #15542 review follow-up
top_k_/top_p_/temperature_sampling_batchpass-through wrappers intotop_k_top_p_sampling_batchvia default arguments (top_k=Nonedisables top-k,top_p=1disables top-p).greedy()wrapper; thegreedyre-export now aliasesgreedy_search_sampling_batchdirectly.Merge flashinfer op variants (ops/flashinfer.py) — #15542 review follow-up
single signature with keyword-only randomness arguments (6 wrappers -> 4).
flashinfer's own API (>= 0.6.4) already accepts both in one signature, with explicit
seed/offsettaking precedence overgenerator; the split wrappers predated that.seed/offset(device tensors) remain requiredfor CUDA-graph capture and cross-rank determinism;
generatorremains for eager paths.The module docstring now documents this distinction.
Remove dead code
trtllm::compute_probs_from_logits_opchain, unused end to end (no Python or C++callers): the
register_fakestub, the thop wrapper +TORCH_LIBRARYregistration,and the exclusive kernel chain in
dynamicTreeKernels.cu(topKProbStage1/2,invokeTopKTopPMaskingForProbs,applyTopKTopPForProbOp,computeSoftmaxForProbOp,computeProbsFromLogits) plus the torch/ATen includes only they used.TorchSamplermethods with zero callers:_handle_finish_reasons,_longest_stop_word_len,_requests_with_stop_words,_request_indices_with_stop_words.Promote de-facto public internals
sampling_utilsandsampler.pywere reaching into underscore-private members ofops.vanilla; promote the actually-shared surface and remove a duplicated constant:_GREEDY_TEMPERATURE_THRESHOLD->GREEDY_TEMPERATURE_THRESHOLD;DISABLE_TEMP_VALin
speculative/interface.pynow derives from it (/ 10, value unchanged at 1e-5)so the sentinel and threshold cannot drift apart silently.
_safely_apply_temperature_inplace->safely_apply_temperature_inplace,_Fusions->Fusions.BEAM_SEARCH_PAD_TOKEN(inops.vanilla, re-exported by thefacade); previously defined independently in two files.
sampling_utilsis now the only module importing theops.*backend modules.Correctness / readability fixes in TorchSampler
_get_logprobs_from_request: use forward slicing for the history view;[:, :-preallocate_extra_steps, :]silently yields an empty view whenpreallocate_extra_steps == 0(the default), skipping history filling._apply_min_length_penalty: returnNoneinstead of echoing the in-place-mutatedlogits; defer the host tensor->list conversions until after the early-out (the common
no-min-length path now skips them); flatten double-nested guard
ifs into earlycontinues, hoisting the per-request offset accumulation above the guards._flashinfer_check_nans: restore the explanatory comments lost in the [TRTLLM-13212][refactor] Centralize sampling logic, split backends into isolated modules #15542 move —the constant
Falsedeliberately keeps flashinfer's syncingcheck_nanpathdisabled while the async device-side assert provides NaN protection without stalling
the pipeline (flashinfer issue Error occured when running medusa inference. #1575).
seed/offsettype hints onsampling_batch_spec_dec_one_model(
Optional[int]->Optional[torch.Tensor]; callers pass device tensors).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.