[TRTLLM-13212][refactor] Centralize sampling logic, split backends into isolated modules#15542
Merged
zhaoyangwang-nvidia merged 10 commits intoJul 8, 2026
Merged
Conversation
zhaoyangwang-nvidia
force-pushed
the
fuse_sampler
branch
from
0206d6e to
9906267
Compare
zhaoyangwang-nvidia
changed the title
zhaoyangwang-nvidia
force-pushed
the
fuse_sampler
branch
6 times, most recently
from
99b8457 to
f21cb8a
Compare
Contributor
📝 WalkthroughWalkthroughConsolidates sampling logic from ChangesSampling Backend Abstraction
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
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/speculative/eagle3_dynamic_tree.py (1)
780-787: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winKeep the
self.use_flashinfergate on this sampling path.sampling_batch_spec_dec_one_model()only checksIS_FLASHINFER_AVAILABLE, while this worker still usesself.use_flashinferto enforce theflashinfer>=0.6.4version gate. Thread that condition through here or reintroduce the check before calling the helper.🤖 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/speculative/eagle3_dynamic_tree.py` around lines 780 - 787, The sampling path in `Eagle3DynamicTree` is missing the `self.use_flashinfer` guard, so it can call `sampling_batch_spec_dec_one_model()` even when the worker should enforce the flashinfer version gate. Update the logic around `sampled` in `eagle3_dynamic_tree.py` to either thread `self.use_flashinfer` into this branch or reintroduce the check before invoking `sampling_batch_spec_dec_one_model()`, keeping the existing flashinfer availability/version behavior consistent.
🤖 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/sampling_backends/backend_torch.py`:
- Around line 460-475: Short-circuit greedy rows in compute_probs_from_logits_op
before calling _apply_temperature() so temperatures at or below
_GREEDY_TEMPERATURE_THRESHOLD never divide logits by zero or a near-zero
sentinel. Use the existing is_greedy check in compute_probs_from_logits_op to
split out those rows, return the one-hot argmax for them immediately, and only
run _apply_temperature() and _apply_top_k_top_p() on non-greedy rows to avoid
inf/nan in mixed batches.
In `@tensorrt_llm/_torch/pyexecutor/sampling_utils.py`:
- Around line 1019-1024: The FlashInfer path in compute_probs_from_logits()
ignores skip_temperature, so make the IS_FLASHINFER_AVAILABLE branch behave like
the custom backend branch by honoring that flag before applying temperatures.
Update the call through backend_flashinfer.compute_probs_from_logits_op in
sampling_utils.py (and any related helper it uses) so the public API produces
consistent probabilities regardless of whether FlashInfer is installed.
- Around line 1040-1046: Add a greedy fast path in the sampling helper before
calling backend_torch._apply_temperature so rows with temperature == 0 or <=
1e-4 are handled via argmax rather than scaled sampling. Update the logic around
sanitize_top_k, backend_torch._apply_temperature, and the flashinfer/native
sampling branch so mixed batches can split greedy rows from stochastic rows
safely and avoid passing inf/nan logits downstream.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 57-69: rejection_sampling_one_model() currently only gates on
IS_FLASHINFER_AVAILABLE, so older FlashInfer installs can still reach
flashinfer.sampling.chain_speculative_sampling and fail later. Update the guard
in rejection_sampling_one_model() (and the
_sample_and_accept_draft_tokens_rejection() path that calls it) to match
SpecWorkerBase.__init__ by checking for FlashInfer 0.6.4+ and raising the same
flashinfer>=0.6.4 RuntimeError before calling chain_speculative_sampling.
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 1514-1520: The bound-backend contract check in the Torch sampler
test is incomplete because it only verifies sample_grouped_strategies, so a
partial bind could still pass. Update the test in test_torch_sampler.py to
assert that sampler._bound_backend matches expected_cls for
sample_grouped_strategies, strategy_grouping_key, and
get_metadata_type_for_group. Use the existing expected_cls selection
(FlashInferGroupedStrategySampler vs SimpleGroupedStrategySampler) and compare
each callable directly so the refactor coverage is complete.
---
Outside diff comments:
In `@tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py`:
- Around line 780-787: The sampling path in `Eagle3DynamicTree` is missing the
`self.use_flashinfer` guard, so it can call
`sampling_batch_spec_dec_one_model()` even when the worker should enforce the
flashinfer version gate. Update the logic around `sampled` in
`eagle3_dynamic_tree.py` to either thread `self.use_flashinfer` into this branch
or reintroduce the check before invoking `sampling_batch_spec_dec_one_model()`,
keeping the existing flashinfer availability/version behavior consistent.
🪄 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: 307a96cf-6f0e-4a31-8ec9-4dda2f274310
📒 Files selected for processing (14)
tensorrt_llm/_torch/pyexecutor/sampler.pytensorrt_llm/_torch/pyexecutor/sampling_backends/__init__.pytensorrt_llm/_torch/pyexecutor/sampling_backends/backend_custom.pytensorrt_llm/_torch/pyexecutor/sampling_backends/backend_flashinfer.pytensorrt_llm/_torch/pyexecutor/sampling_backends/backend_torch.pytensorrt_llm/_torch/pyexecutor/sampling_utils.pytensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.pytensorrt_llm/_torch/speculative/drafting_loops.pytensorrt_llm/_torch/speculative/dynamic_tree_ops.pytensorrt_llm/_torch/speculative/eagle3_dynamic_tree.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/one_model_sampler.pytests/unittest/_torch/sampler/test_logits_logprobs.pytests/unittest/_torch/sampler/test_torch_sampler.py
💤 Files with no reviewable changes (2)
- tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py
- tensorrt_llm/_torch/speculative/one_model_sampler.py
zhaoyangwang-nvidia
force-pushed
the
fuse_sampler
branch
from
f21cb8a to
df07465
Compare
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #56574 [ run ] triggered by Bot. Commit: |
QiJune
reviewed
Jun 30, 2026
QiJune
reviewed
Jun 30, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #56586 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #56574 [ run ] completed with state |
Collaborator
|
PR_Github #56586 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #56788 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #56788 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #57764 [ run ] completed with state |
EmmaQiaoCh
approved these changes
Jul 7, 2026
…to isolated modules - Merge sampling_utils_flashinfer.py into sampling_utils.py (deleted former) - Move _StrategyImpls and FlashInferGroupedStrategySampler to sampling_utils.py - Wrap all flashinfer kernel calls in backend_flashinfer.py wrapper functions - Move _st_* helpers to backend_torch.py (drop prefix, rename to forward_native_sampling etc.) - Move rejection_sampling_one_model to speculative/interface.py (draft acceptance logic) - Remove use_flashinfer param from sampling_batch_spec_dec_one_model (auto-detect) - Add sanitize_top_k (top_k<=0 means keep-all) and apply it in the spec-dec sampling facade - Add torch fallback path to sampling_batch_spec_dec_one_model_for_rejection - Fix top_p renormalization: drop in-place cumsum(out=probs_sorted) that corrupted the denominator - Fix sampler.py import: FlashInferGroupedStrategySampler now from sampling_utils - Move StrategyMetadata to backend_torch.py so BeamSearchMetadata can inherit it - Use explicit re-export form for _Fusions, get_rejected_indices, sample_rejected - Add type annotations to _StrategyImpls subclass sample() overrides and from_strategies - Update sampler unit tests for the new _bound_backend dispatch + module locations Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…_backend Address review feedback (naming consistency + a missed call site): - Rename sampling_backends/ -> sampling_backend/ (singular, matching attention_backend/) and drop the redundant backend_ filename prefix: backend_torch.py -> vanilla.py, backend_flashinfer.py -> flashinfer.py, backend_custom.py -> trtllm.py (the latter wraps torch.ops.trtllm custom ops, so trtllm.py is more explicit than "custom"). - Move SamplerConfig / BoundSamplingBackend / resolve_sampling_backend into sampling_backend/interface.py (mirrors attention_backend/interface.py). - Update all imports, attribute accesses, and docstrings accordingly. - Fix eagle3_dynamic_tree.py: drop the stale use_flashinfer kwarg from the context-subset sampling_batch_spec_dec_one_model call (the parameter was removed earlier in this refactor); the call now matches the new signature. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…dSamplingBackend Address review feedback: TorchSampler only consumes the strategy-grouping callbacks (strategy_grouping_key / get_metadata_type_for_group / sample_grouped_strategies). The compute_probs / sample_from_probs / sample_from_logits / greedy fields bound into BoundSamplingBackend were never read by any call site. - Remove those 4 unused fields from BoundSamplingBackend and resolve_sampling_backend() - Drop the now-dead generic wrappers they pointed at in vanilla.py and flashinfer.py (keep vanilla.greedy, still re-exported for drafting_loops / speculative.interface) - Drop the corresponding aliases / re-exports in sampling_utils.py Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ith kernels/ providers
Address review feedback: the sampling_backend/ name implied flashinfer /
vanilla / trtllm were three peer sampler backends for the upper layer to
select. Reorganize so the abstraction is operation-centric and the
implementation sources stay private.
- New package pyexecutor/sampler/:
sampler.py orchestration (Sampler / TorchSampler / TRTLLMSampler)
sampling_utils.py operation-level facade
kernels/ private provider implementations, selected internally
flashinfer.py, vanilla.py, interface.py
- sampler/__init__.py lazily (PEP 562) re-exports sampler.py symbols so existing
`pyexecutor.sampler` import paths keep working while importing lightweight
submodules (e.g. sampling_utils) avoids pulling in sampler.py's heavy
dependency chain — this also breaks the import cycle with speculative.interface.
- Inline the single TRT-LLM op into the compute_probs_from_logits() facade and
drop the thin trtllm.py wrapper (no standalone provider contract yet).
- Update all import paths across pyexecutor / speculative / auto_deploy / tests.
Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
- BoundSamplingBackend: make it a frozen dataclass parameterised over the strategy-key type (Generic[_StrategyKeyT]) with precise Callable field types instead of Any, so the annotations are actually checked. - Rename the generator-based flashinfer ops to *_generator_op (top_k_sampling_from_probs_generator_op / top_p_sampling_from_probs_generator_op) for naming consistency with the other generator-based ops. - Update pyproject.toml mypy override to cover the whole sampler package (tensorrt_llm._torch.pyexecutor.sampler.*) after the package reorg, so the strict typing established in the base branch is not regressed. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…it, tidy sampler kernels Follow-up on reviewer feedback: - Revert to binding the grouped-sampler class as a unit (self._grouped_sampler_cls) instead of splitting the tightly-coupled grouping callables into separate fields; drop BoundSamplingBackend. resolve_sampling_backend now returns the chosen GroupedStrategySampler subclass and takes is_cuda: bool (redundant IS_FLASHINFER_AVAILABLE check removed at the call site). - Rename _safely_apply_temperature -> _safely_apply_temperature_inplace (arg logits_inout) to make the in-place contract explicit. - vanilla.compute_probs_from_logits_op: write greedy rows into probs in place, removing the full [batch, vocab] one-hot buffer and torch.where copy (numerically equivalent; kept dense to stay torch.compile-safe). - Move flashinfer op imports to the top of sampling_utils.py (unconditional); fix sampling_batch_spec_dec_one_model return annotation (tuple -> Tensor). - drafting_loops: unpack greedy() via tokens, _ = ... - Update sampler unit tests for the class-based _grouped_sampler_cls dispatch. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…mpler package The sampler refactor moved sampler.py, sampling_utils.py and sampling_utils_flashinfer.py into the sampler/ package, so the old paths in static-analysis-files no longer existed and mypy silently skipped the new files. Point the list at the new module paths. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Align the sampler subpackage with the existing codebase convention for provider-specific compute implementations (cf. _torch/modules/fused_moe/ops), where 'ops' denotes swappable backend implementations behind an interface, reserving 'kernels' for actual CUDA/C++ kernel sources. No behavior change; updates all imports, the lazy __getattr__ submodule guard, the mypy static-analysis file list, and doc comments. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ing type-check Enabling static analysis on the moved sampler files (previous commit pointed static-analysis-files at the real module paths) surfaced pre-existing [no-any-return] errors: functions declared to return Tensor were returning the Any-typed results of untyped flashinfer / torch.ops calls. Bind the results to annotated locals before returning, matching the existing convention. No behavior change. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
fuse_sampler
branch
from
4823b81 to
5c0b3ec
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57973 [ run ] triggered by Bot. Commit: |
chzblych
approved these changes
Jul 7, 2026
Collaborator
|
PR_Github #57973 [ run ] completed with state
|
Collaborator
|
/bot run |
Collaborator
|
PR_Github #58118 [ run ] triggered by Bot. Commit: |
sunnyqgg
approved these changes
Jul 8, 2026
Collaborator
|
PR_Github #58118 [ run ] completed with state
|
… in spec sampling Address review on the one-model speculative sampling helpers: - sampling_batch_spec_dec_one_model: cast the greedy argmax tokens to the sampler's dtype before torch.where so the result stays int32 instead of being promoted to int64 (flashinfer returns int32). - sampling_batch_spec_dec_one_model_for_rejection: restore the pre-refactor RuntimeError when flashinfer is missing instead of silently falling back to a global-RNG torch sample that ignores seed/offset, which would break determinism and cross-rank consistency. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #58136 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58136 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58167 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58167 [ run ] completed with state |
BrianLi23
pushed a commit
to BrianLi23/TensorRT-LLM
that referenced
this pull request
Jul 9, 2026
…to isolated modules (NVIDIA#15542) Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
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
8 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
Description
Refactors the PyTorch backend sampling subsystem into a single public interface backed by separate backend implementation modules.
Before: Logic split across two overlapping stacks:
sampling_utils.py+sampling_utils_flashinfer.py— host-side strategy samplingspeculative/one_model_sampler.py— one-model spec-dec sampling (a separate stack)After: One interface + three backend implementations:
pyexecutor/
├── sampling_utils.py ← single public interface
└── sampling_backends/
├── backend_torch.py ← pure PyTorch kernels
├── backend_flashinfer.py ← FlashInfer kernel wrappers
└── backend_custom.py ← C++ / custom op wrappers
This "interface + backend implementations" structure is designed to pave the way for future performance optimization and feature extension:
sampling_utils.pyinterface; which backend a kernel lives in and how it is implemented stays transparent to the caller. Whether we later replace, tune, or add kernels, the changes are confined to the backend implementation modules and never leak to the call sites.backend_custom.pyis the landing spot for new custom kernels: when we want to support new sampling capabilities for the sampler, or do kernel-level optimization for a specific scenario, we can drop the new custom kernel (C++ op / Triton / hand-written CUDA, etc.) straight intobackend_custom.pyand wire it into dispatch viaresolve_sampling_backend()— without touching the interface API or the other backends.resolve_sampling_backend()into aBoundSamplingBackend, so there is no per-call backend selection at runtime. New custom kernels fall naturally into the same binding mechanism, making them easy to reuse inside the CUDA graph capture path.Key Changes
sampling_utilsor each other, preventing circular dependenciessampling_utils_flashinfer.pydeleted:_StrategyImplsandFlashInferGroupedStrategySamplermerged directly intosampling_utils.pyone_model_sampler.pydeleted: sampling functions migrated tosampling_utils.py;rejection_sampling_one_modelmoved tospeculative/interface.py(draft-token acceptance logic belongs on the speculative side)BoundSamplingBackendbinds all kernel functions at init time viaresolve_sampling_backend(); no per-call backend selection inside CUDA graph captureuse_flashinferparam removed fromsampling_batch_spec_dec_one_model(auto-detects viaIS_FLASHINFER_AVAILABLE)sanitize_top_k: centralized top-k normalization in the spec-dec interface —top_k <= 0(and oversized disable sentinels) map tovocab_size(keep-all), mirroring the C++ op, instead of breaking the flashinfer / torch top-k paths on a literal0sampling_batch_spec_dec_one_model_for_rejectionnow supports a pure-torch path when FlashInfer is unavailabledrafting_loops.pyanddynamic_tree_ops.pynow import directly fromsampling_utilsinstead ofone_model_samplerBehavior fix included in this refactor
torch.cumsum(probs_sorted, dim=-1, out=probs_sorted), aliasing the cumulative-probability tensor that is reused as the renormalization denominator. The subsequentmasked_fill_(...)corrupted that denominator, soprobsno longer summed to 1. Fixed by removing the in-placeout=socumulative_probsis allocated separately (two call sites inbackend_torch.py).Files Changed
sampling_backends/__init__.py,backend_torch.py,backend_flashinfer.py,backend_custom.pysampling_utils.py,sampler.py,speculative/interface.py,speculative/drafting_loops.py,speculative/dynamic_tree_ops.py,speculative/eagle3_dynamic_tree.pytests/unittest/_torch/sampler/test_torch_sampler.py,test_logits_logprobs.py— updated for the new_bound_backenddispatch and module locationssampling_utils_flashinfer.py,speculative/one_model_sampler.pyTesting
tests/unittest/_torch/sampler/test_torch_sampler.pypasses (262 passed, 162 skipped) on B200 against amain-matched build.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.