[TRTLLM-13231][feat] Support top_p_decay in the PyTorch TorchSampler

[zhaoyangwang-nvidia]

Implement Top-P Decay on the FlashInfer sampling path, porting the semantics of the legacy C++ computeToppDecay kernel. The three SamplingParams fields top_p_decay / top_p_min / top_p_reset_ids were previously dropped by the PyTorch backend; they are now plumbed through and honored.

Each decode step, after the token is sampled, the per-request runtime top-p is updated as runtime_p = max(runtime_p * top_p_decay, top_p_min), or reset to its initial value when the sampled token equals top_p_reset_ids (gated on reset_id >= 0). The runtime top-p is carried as a per-row tensor into the FlashInfer batched sampling ops, never encoded into the cached strategy tuple.

Scope is limited to the FlashInfer path and single-token decode steps; decay combined with speculative decoding, beam search, or the non-FlashInfer backend is rejected early. Out-of-range top_p_decay / top_p_min warn-and-default, matching the C++ runtime.

Adds unit tests covering activation/routing, warn-and-default validation, numeric parity with the C++ recurrence, slot-reuse lifecycle, and mixed decay / non-decay batches.

@coderabbitai summary

Description

Implements Top-P Decay in the PyTorch TorchSampler, porting the semantics of
the legacy C++ computeToppDecay kernel. The three SamplingParams fields
(top_p_decay / top_p_min / top_p_reset_ids) already reached the C++/legacy
path but were silently dropped by the PyTorch backend — they are now plumbed
through and honored end-to-end.

Each decode step, after a token is sampled, the per-request runtime top-p is
updated as runtime_p = max(runtime_p * top_p_decay, top_p_min), or reset to its
initial value when the sampled token equals top_p_reset_ids (gated on
reset_id >= 0). The runtime top-p is carried as a per-row tensor into the
FlashInfer batched sampling ops and is never encoded into the cached strategy
tuple, so the same value drives both sampling and processed logprobs.

Scope

• FlashInfer path only, single-token decode steps only.
• Decay combined with speculative decoding (req_num_steps > 1), beam search, or
  the non-FlashInfer backend is rejected early with a clear error.
• An explicit greedy control (top_k == 1 / top_p == 0.0 / temperature == 0)
  takes precedence over decay; out-of-range top_p_decay / top_p_min warn and
  fall back to defaults (C++ parity). top_p_min > top_p is allowed.

Implementation

• sampling_params.py — warn-and-default validation; _greedy_decoding made
  consistent with decay routing; corrected top_p_reset_ids docstring default.
• sampling_utils.py — UtilsSamplingParams + activation predicate, decay-aware
  strategy routing, per-row runtime-top-p override.
• sampler.py — per-slot CUDA TopPDecayStore + lifecycle gate, early
  rejection, per-step runtime-top-p sourcing, post-sample decay/reset update.

Test Coverage

New tests/unittest/_torch/sampler/test_top_p_decay.py (30 tests): activation /
strategy routing, warn-and-default validation, numeric parity against a Python
port of the C++ computeToppDecay recurrence (multi-step + reset + min>initial),
slot-reuse lifecycle, mixed decay / non-decay batches, and rejection paths.

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.