[TRTLLM-14736][chore] Split the sampler package into per-feature modules

[zhaoyangwang-nvidia]

Dev Engineer Review

• Refactored the sampler package into focused per-feature modules:
  • finish_reasons.py for finish-reason handling (FinishReasonsHandler).
  • top_p_decay.py for Top-P decay lifecycle management (TopPDecayHandler/TopPDecayStore).
  • logprobs.py for log-prob state, storage, and tensor↔list conversion helpers.
  • sampler_common.py for shared sampler utilities and small request/tensor helpers (UtilsSamplingParams, index defaults, token append helper).
  • sampler_strategy.py replacing sampling_utils.py and introducing request→strategy routing (_request_strategy).
• Moved speculative-decoding sampling operations into sampler/ops/* modules:
  • Added/extended flashinfer-based speculative utilities in ops/flashinfer.py (including sanitize_top_k and speculative one-model helpers with per-request tensors).
  • Updated eager/vanilla call sites to import from sampler/ops/vanilla.py.
• Preserved external import compatibility via lazy package forwarding in sampler/__init__.py, while removing/redirecting direct imports of the old sampling_utils module.
• Updated call sites across Torch executor and speculative modules to use the new module boundaries (including routing changes for strategy resolution, logprobs handling, finish-reasons handling, and Top-P decay).
• Reduced sampler.py surface area (per author notes) and re-established an acyclic sampler dependency graph with maximum depth four.
• Verified test-list/config/waive files were not changed (no diffs under tests/integration/test_lists/, test-db/, qa/, or waives.txt).

QA Engineer Review

Modified test code:

• tests/unittest/_torch/sampler/test_beam_search.py
  • Updated beam-search logprob conversion to use convert_logprobs_tensor_to_list from tensorrt_llm._torch.pyexecutor.sampler.logprobs.
  • Affected coverage exercised by existing tests including:
    • test_beam_search_e2e, test_beam_search_disagg_e2e
    • test_logprobs_trtllm_sampler, test_logprobs_torch_sampler
    • plus other beam-history/finish-related tests in the file (test_create_beam_history, test_finish_beams, etc.).
• tests/unittest/_torch/sampler/test_logits_logprobs.py
  • Updated _StrategyImpls import to tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy.
  • Affects existing logprob suite tests such as test_logprobs_simple_format, test_processed_logprobs_e2e, and test_logprobs_with_grouped_samplings_strategies.
• tests/unittest/_torch/sampler/test_torch_sampler.py
  • Updated FinishReasonsHandler monkeypatch targets to match the new import location.
  • Updated Top-P decay test runtime update parity to use sampler._top_p_decay new internal API (store, _slots, update_after_sample).
  • Affects existing tests/classes including TestFinishReasons methods (test_are_stop_words_*, stop-word buffer resize tests) and TestTopPDecay.test_runtime_update_parity.

Reported validation: sampler and token-ban suites passed (246 passed, 0 failed).

Verdict: needs follow-up (no CBTS/CI test-list coverage metadata provided for the modified test modules).

Description

sampler.py had grown to ~5.8k lines, with TorchSampler alone accounting for
~3.9k of them. This PR splits it into focused modules and gives the package a
dependency-free base layer. It is a pure refactor: no behavior change.

Resulting structure

                         sampler.py  (4319)
                        orchestration
                              │
       ┌──────────┬───────────┼────────────┬──────────────┐
       ▼          ▼           ▼            ▼              ▼
  token_ban   logprobs  finish_reasons  top_p_decay  sampler_strategy
    (708)      (275)        (873)          (377)         (905)
       │          │           │             │  │            │
       │          │           │             │  └──► ops/vanilla ◄──┐
       │          │           │             │                      │
       └──────────┴───────────┴─────────────┴──────────────────────┤
                              ▼                             ops/flashinfer
                      sampler_common  (154)
                  no intra-package imports

module	depends on (intra-package)
sampler_common	—
logprobs	—
token_ban	—
ops/vanilla	—
ops/flashinfer	ops/vanilla
finish_reasons	sampler_common
sampler_strategy	sampler_common, ops/*
top_p_decay	sampler_common, sampler_strategy, ops/vanilla
sampler	all of the above

Acyclic, max depth 4.

What moved

finish_reasons.py — FinishReasonsHandler was an 847-line class nested
inside TorchSampler that never touched the outer instance: it takes plain
sizes at construction and owns all of its own state (the finish-reason store,
the stop-word rule and past-token buffers with their resize/refresh logic, and
the end-id / max-length / stop-word device checks). Promoted to a module-level
class.

top_p_decay.py — the Top-P Decay store and its full lifecycle (admission,
pre-sample metadata, post-sample recurrence, retirement), now behind a
TopPDecayHandler. The handler owns both the per-slot buffers and the
host-side membership set that gates them, so the two can no longer drift —
previously the set was a TorchSampler attribute while the store hung off
TorchSampler.Store, even though the set is the sole gate for reading the
store.

logprobs.py — the log-probs data types (LogProbsState,
LogProbsStateList, LogProbsStore) and the three conversion helpers that
took no sampler state, now module-level functions. Log-probs sizing
(_prepare_log_probs) and the per-step device gather (_process_logprobs)
deliberately stay in TorchSampler: they read and grow the sampler-wide shapes
(TOPK_LOGPROBS_SHAPE and friends), so they are bound to sampler state rather
than to the log-probs data.

sampler_common.py — the package's base layer. Holds everything a module
can need without resolving a sampling strategy: the per-request queries
(_request_get_sampling_params and friends), UtilsSamplingParams, the shared
DEFAULT_BEAM_IDX / DEFAULT_STEP_IDX constants, FinishReasonsList, and the
tensor helpers. It imports nothing from the package, so any module may depend
on it.

UtilsSamplingParams lives here rather than in sampler_strategy on purpose:
it is the plain scalar view of a request's sampling config — the input to
strategy resolution, not a product of it — and several features read it without
ever resolving a strategy.

sampler_strategy.py (renamed from sampling_utils.py) — strategy
resolution, the strategy implementations and the grouped samplers. Keeps
_request_strategy, the one request-facing helper that has to know what a
Strategy is, and reads its inputs from sampler_common.

Also in this branch: the spec-decoding sampling ops move out of
sampling_utils into ops/flashinfer.py (they hard-depend on flashinfer),
decoupling the speculative and auto_deploy call sites.

Why the base layer matters beyond this PR

Two in-flight PRs hit circular imports because per-request helpers lived in
sampler.py:

• [TRTLLM-13229][feat] implement repetition / frequency / presence penalties for TorchSampler #16485 (repetition / frequency / presence penalties) — review feedback asked
  for all penalty code in one file under pyexecutor/sampler/, which was
  blocked on exactly this.
• [TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, diversity_rate, early_stopping, VBWS, and CBA performance #16620 (beam search length_penalty, diversity_rate, early_stopping).

Neither needs Strategy. The penalty code uses _unwrap_singleton,
_get_max_beam_width and DEFAULT_BEAM_IDX and never resolves a strategy —
all three now sit in sampler_common, so a penalties.py can join the feature
layer next to token_ban.py without importing sampler.py.

sampler.py drops from 5812 to 4319 lines. External import paths are
unaffected — everything still resolves through the package's lazy __getattr__
forwarding, so pyexecutor.sampler.TorchSampler, .DEFAULT_BEAM_IDX etc. keep
working.

Test Coverage

tests/unittest/_torch/sampler/test_torch_sampler.py and test_token_ban.py:
246 passed, 0 failed (H200). These cover all three extracted feature areas —
TestFinishReasons (whose monkeypatch targets moved to the promoted class), the
top-p decay runtime-update parity test (whose store/membership access paths
changed), and the token-ban paths.

Test-only updates: monkeypatch targets move from
TorchSampler.FinishReasonsHandler to the imported class, the top-p decay test
reaches the store via the handler, and test_beam_search.py imports
convert_logprobs_tensor_to_list from logprobs.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[zhaoyangwang-nvidia]

Hi @lori-ren please help to review, let's merge this PR ASAP.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1960dec0-8989-4028-81ce-1de3ffea57f8

📥 Commits

Reviewing files that changed from the base of the PR and between d4b3d45 and 49c7b29.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
• tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py
• tensorrt_llm/sampling_params.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py

---

Walkthrough

TorchSampler functionality is reorganized into dedicated finish-reason, logprob, top-p decay, request-parameter, and sampling-operation modules. Speculative decoding imports move to vanilla and FlashInfer operation modules, with corresponding sampler wiring and test updates.

Changes

Torch sampler modularization

Layer / File(s)	Summary
Finish-reason and logprob handlers tensorrt_llm/_torch/pyexecutor/sampler/{finish_reasons,logprobs}.py, tensorrt_llm/_torch/pyexecutor/sampler/sampler.py	Finish-reason tracking and logprob conversion/storage are extracted into dedicated modules and integrated into TorchSampler.
Top-P decay lifecycle tensorrt_llm/_torch/pyexecutor/sampler/{top_p_decay,sampler}.py, tests/unittest/_torch/sampler/test_torch_sampler.py	Top-P decay state, validation, metadata, updates, retirement, and related test access move into TopPDecayHandler.
Request strategy resolution tensorrt_llm/_torch/pyexecutor/sampler/{sampler_common,sampler_strategy,__init__}.py, tensorrt_llm/_torch/pyexecutor/llm_request.py	Request sampling parameters and strategy caching are separated into shared helpers and sampler_strategy.py, with lazy routing updated.
Speculative sampling operations tensorrt_llm/_torch/pyexecutor/sampler/ops/{vanilla,flashinfer}.py, tensorrt_llm/_torch/speculative/*	Speculative probability and sampling functions are provided by FlashInfer and vanilla operation modules, with callers updated to use them.
Integration and tests tensorrt_llm/_torch/auto_deploy/shim/demollm.py, tests/unittest/_torch/sampler/*, tensorrt_llm/sampling_params.py	Imports, private access paths, monkeypatch targets, and documentation references are updated for the new module boundaries.

Estimated code review effort: 5 (Critical) | ~90+ minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#8132: Refactors sampling strategy selection.
• NVIDIA/TensorRT-LLM#8794: Reworks finish-reason and stop-criteria handling.
• NVIDIA/TensorRT-LLM#16365: Routes sampling through vanilla and FlashInfer operation APIs.

Suggested labels: api-compatible

Suggested reviewers: qijune

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and accurately summarizes the main refactor into per-feature sampler modules.
Description check	✅ Passed	The description follows the template well, with clear problem/solution, test coverage, and checklist sections filled in.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (4)

tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py (1)

219-244: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

@dataclass(kw_only=True) is inert here.

_TemporaryData declares no fields and defines its own __init__, so the decorator contributes nothing (and its generated __init__ is shadowed). Either drop the decorator or express the state as dataclass fields with default_factory and let clear() reassign.

♻️ Field-based alternative

-    `@dataclass`(kw_only=True)
     class _TemporaryData:
         """Data structure to store the temporary data during setup_sampler_step for new requests"""
 
         def __init__(self) -> None:

🤖 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/finish_reasons.py` around lines 219 -
244, Remove the inert `@dataclass`(kw_only=True) decorator from _TemporaryData,
since the class defines no dataclass fields and provides its own __init__.
Preserve the existing initialization and clear() behavior unchanged.

tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py (2)

110-120: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Avoid tensor allocation to compute a scalar max.

_get_max_beam_width builds a torch.Tensor from beam_width_array just to call .max().item(). Since this path is not cached for beam-search requests (_request_sampling_params_cachable returns False when use_beam_search), and _request_strategy is invoked multiple times per sampling step, this tensor allocation runs repeatedly on the hot path for every beam-search request. Plain Python max() on the list avoids the allocation entirely.

⚡ Proposed fix

     if sampling_config.beam_width_array is not None:
-        max_beam_width = max(
-            max_beam_width,
-            cast(
-                int, torch.tensor(sampling_config.beam_width_array, dtype=torch.int32).max().item()
-            ),
-        )
+        max_beam_width = max(max_beam_width, max(sampling_config.beam_width_array))

🤖 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_common.py` around lines 110 -
120, Update _get_max_beam_width to compute the maximum of
sampling_config.beam_width_array with Python’s built-in max() instead of
constructing a torch.Tensor and calling tensor.max().item(), while preserving
the existing comparison with beam_width and None handling.

Source: Learnings

---

31-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add __all__ for the public surface.

This module exports several public names (UtilsSamplingParams, int_tensor, add_token, DEFAULT_BEAM_IDX, DEFAULT_STEP_IDX, FinishReasonsList) but has no __all__, unlike the sibling sampler_strategy.py which defines one for its public surface.

As per coding guidelines: "keep __all__ updated for public interfaces."

🤖 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_common.py` around lines 31 -
46, Add an __all__ declaration in sampler_common.py listing the module’s public
interfaces: UtilsSamplingParams, int_tensor, add_token, DEFAULT_BEAM_IDX,
DEFAULT_STEP_IDX, and FinishReasonsList. Keep the list synchronized with these
existing exported symbols and follow the pattern used by sampler_strategy.py.

Source: Coding guidelines

tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py (1)

283-303: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the in-place logits mutation.

sampling_batch_spec_dec_one_model mutates the caller's logits tensor in place via vanilla.safely_apply_temperature_inplace (a div_). Current callers pass disjoint slices so this is currently safe, but the docstring doesn't disclose the mutation, which is risky for a shared utility with multiple call sites across speculative-decoding backends.

📝 Proposed docstring update

     """CUDA-graph compatible sampling; supports mixed sampling params. Returns sampled tokens."""
+    # NOTE: mutates `logits` in place (via safely_apply_temperature_inplace).
+    # Callers must not reuse the same tensor/view afterwards.
     top_k = sanitize_top_k(top_k, logits.shape[-1])

🤖 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/flashinfer.py` around lines 283 -
303, Update the docstring of sampling_batch_spec_dec_one_model to explicitly
state that it mutates the input logits tensor in place during temperature
application; leave the sampling behavior and implementation unchanged.

🤖 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/ops/flashinfer.py`:
- Around line 265-280: Update compute_probs_from_logits so top_k and top_p
always remain tensor inputs during the compiled path, replacing None-based
toggles with the established sentinel or mask representation expected by
compute_probs_from_logits_op. Preserve sanitization and ensure
skip_top_k/skip_top_p changes do not create separate torch.compile
specializations.

In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 1394-1396: Move the _top_p_decay initialization into the
torch.inference_mode(False) block alongside self.store and the other mutable
sampler state. Keep TopPDecayHandler construction unchanged otherwise, ensuring
setup_sampler_step() operates on regular mutable tensors rather than inference
tensors.

---

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py`:
- Around line 219-244: Remove the inert `@dataclass`(kw_only=True) decorator from
_TemporaryData, since the class defines no dataclass fields and provides its own
__init__. Preserve the existing initialization and clear() behavior unchanged.

In `@tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py`:
- Around line 283-303: Update the docstring of sampling_batch_spec_dec_one_model
to explicitly state that it mutates the input logits tensor in place during
temperature application; leave the sampling behavior and implementation
unchanged.

In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py`:
- Around line 110-120: Update _get_max_beam_width to compute the maximum of
sampling_config.beam_width_array with Python’s built-in max() instead of
constructing a torch.Tensor and calling tensor.max().item(), while preserving
the existing comparison with beam_width and None handling.
- Around line 31-46: Add an __all__ declaration in sampler_common.py listing the
module’s public interfaces: UtilsSamplingParams, int_tensor, add_token,
DEFAULT_BEAM_IDX, DEFAULT_STEP_IDX, and FinishReasonsList. Keep the list
synchronized with these existing exported symbols and follow the pattern used by
sampler_strategy.py.

🪄 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: db75bfe6-ad85-42cd-ba2f-bb070ca7d469

📥 Commits

Reviewing files that changed from the base of the PR and between 15dbae9 and d4b3d45.

📒 Files selected for processing (18)

• tensorrt_llm/_torch/auto_deploy/shim/demollm.py
• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/pyexecutor/sampler/__init__.py
• tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py
• tensorrt_llm/_torch/pyexecutor/sampler/logprobs.py
• tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
• tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py
• tensorrt_llm/_torch/pyexecutor/sampler/top_p_decay.py
• tensorrt_llm/_torch/speculative/drafting_loops.py
• tensorrt_llm/_torch/speculative/dynamic_tree_ops.py
• tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py
• tensorrt_llm/_torch/speculative/interface.py
• tests/unittest/_torch/sampler/test_beam_search.py
• tests/unittest/_torch/sampler/test_logits_logprobs.py
• tests/unittest/_torch/sampler/test_torch_sampler.py

[zhaoyangwang-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62448 [ run ] triggered by Bot. Commit: 49c7b29 Link to invocation

[tensorrt-cicd]

PR_Github #62448 [ run ] completed with state SUCCESS. Commit: 49c7b29
/LLM/main/L0_MergeRequest_PR pipeline #50603 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[zhaoyangwang-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #62508 [ run ] triggered by Bot. Commit: 49c7b29 Link to invocation

[tensorrt-cicd]

PR_Github #62508 [ run ] completed with state SUCCESS. Commit: 49c7b29
/LLM/main/L0_MergeRequest_PR pipeline #50656 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation