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

docs/source/features/sampling.md

@@ -75,7 +75,8 @@ llm.generate(["Hello, my name is",
 
 * The sampling is controlled via `SamplingParams`.
 
-* By default (`temperature = top_p = top_k = None`), greedy sampling is used.
+* By default (`temperature = top_p = top_k = None`), greedy sampling is used
+  (unless top-p decay is active, see below).
 
 * If either `temperature = 0`, `top_p = 0`, and/or `top_k = 1`, is specified, sampling is greedy,
   irrespective of the values of the remaining parameters.
@@ -101,6 +102,19 @@ llm.generate(["Hello, my name is",
 
   * The implementation does not guarantee any particular treatment of tied probabilities.
 
+* Top-P decay is supported: if `top_p_decay < 1` is specified, the effective `top_p` is
+  multiplied by `top_p_decay` after every sampled token, bounded from below by `top_p_min`
+  (default `1e-6`), and reset to the initial `top_p` whenever the token `top_p_reset_ids`
+  is sampled (default `-1`, which never matches a token). Out-of-range values
+  (`top_p_decay` or `top_p_min` outside `(0, 1]`, negative `top_p_reset_ids`) are rejected.
+
+  * An active top-p decay implies top-p sampling even if `top_p` is unspecified or `top_p = 1`
+    (the initial `top_p` then defaults to 1). However, explicitly requested greedy sampling
+    (`temperature = 0`, `top_p = 0`, and/or `top_k = 1`) takes precedence over top-p decay.
+
+  * Top-P decay is not supported in combination with beam search or with speculative decoding
+    modes that route draft tokens through the Torch Sampler; such requests are rejected.
+
 ### Performance
 
 The Torch Sampler leverages the optimized sampling kernels provided by

tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py

@@ -373,3 +373,105 @@ def gather_log_softmax(inputs_cuda: torch.Tensor, indices_cuda: torch.Tensor) ->
         torch._dynamo.mark_dynamic(inputs_cuda, 0)
         torch._dynamo.mark_dynamic(indices_cuda, 0)
         return Fusions._gather_log_softmax_impl(inputs_cuda, indices_cuda)
+
+    # --- Top-P Decay ops ---------------------------------------------------
+    # Host-launch-bound per-step ops (a few dozen elements per row), fused with
+    # Inductor to keep the launch count low. mode="max-autotune-no-cudagraphs":
+    # cudagraphs is unsafe here (the update mutates persistent per-slot state
+    # in place and the gather's output is consumed outside the compiled region;
+    # cudagraph static output buffers get overwritten by subsequent replays).
+    # mark_dynamic on the batch-varying dims avoids recompilation as the batch
+    # composition changes. Compilation is lazy: the first decay-active request
+    # pays it (roughly a second); non-decay workloads never trigger it. See
+    # TorchSampler.TopPDecayStore for the feature-level semantics.
+
+    @staticmethod
+    @torch.compile(mode="max-autotune-no-cudagraphs")
+    def _top_p_decay_update_impl(
+        runtime_top_p: torch.Tensor,
+        initial_top_p: torch.Tensor,
+        top_p_decay: torch.Tensor,
+        top_p_min: torch.Tensor,
+        reset_ids: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        step_tokens: torch.Tensor,
+        sampled_slots: torch.Tensor,
+    ) -> None:
+        active = is_decay_slot[sampled_slots]
+        current = runtime_top_p[sampled_slots]
+        updated = torch.where(
+            step_tokens[sampled_slots] == reset_ids[sampled_slots],
+            initial_top_p[sampled_slots],
+            torch.maximum(current * top_p_decay[sampled_slots], top_p_min[sampled_slots]),
+        )
+        runtime_top_p[sampled_slots] = torch.where(active, updated, current)
+
+    @staticmethod
+    def top_p_decay_update(
+        *,
+        runtime_top_p: torch.Tensor,
+        initial_top_p: torch.Tensor,
+        top_p_decay: torch.Tensor,
+        top_p_min: torch.Tensor,
+        reset_ids: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        step_tokens: torch.Tensor,
+        sampled_slots: torch.Tensor,
+    ) -> None:
+        """Fused in-place update of ``runtime_top_p`` for the sampled decay slots.
+
+        Applies the Top-P Decay recurrence (see ``TorchSampler.TopPDecayStore``
+        for the feature-level semantics) to every sampled row whose slot is
+        decay-active per ``is_decay_slot``.
+
+        All per-slot tensors are 1-D of length ``max_num_sequences``;
+        ``step_tokens`` is a slot-indexed 1-D (possibly strided) int32 view of
+        the new-tokens buffer for a fixed step/beam
+        (``new_tokens[step, :, beam]``); ``sampled_slots`` is 1-D of length
+        ``num_sampled`` (this iteration's rows). ``runtime_top_p`` is mutated
+        in place; nothing is returned.
+        """
+        torch._dynamo.mark_dynamic(sampled_slots, 0)
+        Fusions._top_p_decay_update_impl(
+            runtime_top_p,
+            initial_top_p,
+            top_p_decay,
+            top_p_min,
+            reset_ids,
+            is_decay_slot,
+            step_tokens,
+            sampled_slots,
+        )
+
+    @staticmethod
+    @torch.compile(mode="max-autotune-no-cudagraphs")
+    def _top_p_decay_gather_impl(
+        runtime_top_p: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        static_top_p: torch.Tensor,
+        slots: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.where(is_decay_slot[slots], runtime_top_p[slots], static_top_p)
+
+    @staticmethod
+    def top_p_decay_gather(
+        *,
+        runtime_top_p: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        static_top_p: torch.Tensor,
+        slots: torch.Tensor,
+    ) -> torch.Tensor:
+        """Fused pre-sample per-row top-p gather for decay-active rows.
+
+        Returns a new per-row tensor::
+
+            row_top_p[i] = runtime_top_p[slots[i]]  if is_decay_slot[slots[i]]
+                         = static_top_p[i]          otherwise
+
+        ``runtime_top_p`` / ``is_decay_slot`` are per-slot arrays;
+        ``static_top_p`` and ``slots`` are per-row (length = the group's
+        per-step row count).
+        """
+        torch._dynamo.mark_dynamic(slots, 0)
+        torch._dynamo.mark_dynamic(static_top_p, 0)
+        return Fusions._top_p_decay_gather_impl(runtime_top_p, is_decay_slot, static_top_p, slots)