From 5a0e33f43b8405e9a7984e6c58971dc904d8dee4 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 15 Jul 2026 20:15:42 -0700 Subject: [PATCH 1/9] [TRTLLM-13231][feat] Support top_p_decay in TorchSampler Add Top-P Decay (top_p_decay / top_p_min / top_p_reset_ids) support to the PyTorch backend TorchSampler, with semantics matching the C++ computeToppDecay kernel: after each sampled token, runtime_top_p = initial_top_p if token == reset_id (reset_id >= 0) = max(runtime_top_p * top_p_decay, top_p_min) otherwise - Per-slot runtime state in resident CUDA buffers, gated on-device by an is_decay_slot mask plus a host-side slot set for O(1) early-outs; the hot path needs no host/device sync. - Two fused custom ops (thop/samplerOp.cpp, kernels/samplerKernels/): top_p_decay_gather (pre-sample per-row top-p, replaces index_select x2 + where) and top_p_decay_update (post-sample in-place decay with in-kernel token gather from the strided new_tokens view). - Active decay forces a top-p-capable strategy even when top_p is 1.0/unset; an explicit greedy control (top_k==1 / top_p==0 / temperature==0) wins over decay (single-source predicate SamplingParams.params_imply_explicit_greedy). - Out-of-range decay params follow the C++ warn-and-default policy. - Unsupported combinations (beam search, speculative draft tokens through TorchSampler) are rejected per-request at admission via validate_request, with a req_num_steps == 1 guard at sample time; finished requests retire their slot so the early-outs re-arm. - Unit tests (test_torch_sampler.py::TestTopPDecay): strategy routing, runtime-update parity with the C++ reference (cases ported from topPSamplingLayerTest.cpp), and per-request rejection of unsupported combinations. Signed-off-by: ZhaoyangWang --- .../samplerKernels/toppDecayKernels.cu | 99 +++++ .../kernels/samplerKernels/toppDecayKernels.h | 59 +++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/samplerOp.cpp | 137 +++++++ docs/source/features/sampling.md | 15 +- .../_torch/pyexecutor/sampler/ops/trtllm.py | 78 ++++ .../_torch/pyexecutor/sampler/sampler.py | 352 ++++++++++++++++++ .../pyexecutor/sampler/sampling_utils.py | 93 ++++- tensorrt_llm/sampling_params.py | 49 ++- .../_torch/sampler/test_torch_sampler.py | 114 ++++++ 10 files changed, 992 insertions(+), 5 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu create mode 100644 cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h create mode 100644 cpp/tensorrt_llm/thop/samplerOp.cpp create mode 100644 tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py diff --git a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu new file mode 100644 index 000000000000..addcd038a33e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +__global__ void toppDecayUpdateKernel(float* __restrict__ runtimeTopP, float const* __restrict__ initialTopP, + float const* __restrict__ topPDecay, float const* __restrict__ topPMin, int32_t const* __restrict__ resetIds, + bool const* __restrict__ isDecaySlot, int32_t const* __restrict__ stepTokens, int64_t stepTokenStride, + int64_t const* __restrict__ sampledSlots, int32_t numSampled) +{ + int32_t const i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= numSampled) + { + return; + } + int64_t const slot = sampledSlots[i]; + if (!isDecaySlot[slot]) // on-device decay gate + { + return; + } + int32_t const tok = stepTokens[slot * stepTokenStride]; // gather in-kernel (strided new_tokens view) + int32_t const rid = resetIds[slot]; + float updated; + if (rid >= 0 && tok == rid) + { + updated = initialTopP[slot]; + } + else + { + updated = fmaxf(runtimeTopP[slot] * topPDecay[slot], topPMin[slot]); + } + runtimeTopP[slot] = updated; +} + +__global__ void toppDecayGatherKernel(float* __restrict__ rowTopP, float const* __restrict__ runtimeTopP, + bool const* __restrict__ isDecaySlot, float const* __restrict__ staticTopP, int64_t const* __restrict__ slots, + int32_t numRows) +{ + int32_t const i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= numRows) + { + return; + } + int64_t const slot = slots[i]; + rowTopP[i] = isDecaySlot[slot] ? runtimeTopP[slot] : staticTopP[i]; +} + +} // namespace + +void invokeToppDecayUpdate(float* runtimeTopP, float const* initialTopP, float const* topPDecay, float const* topPMin, + int32_t const* resetIds, bool const* isDecaySlot, int32_t const* stepTokens, int64_t stepTokenStride, + int64_t const* sampledSlots, int32_t numSampled, cudaStream_t stream) +{ + if (numSampled == 0) + { + return; + } + constexpr int32_t kBlock = 256; + int32_t const grid = (numSampled + kBlock - 1) / kBlock; + toppDecayUpdateKernel<<>>(runtimeTopP, initialTopP, topPDecay, topPMin, resetIds, + isDecaySlot, stepTokens, stepTokenStride, sampledSlots, numSampled); +} + +void invokeToppDecayGather(float* rowTopP, float const* runtimeTopP, bool const* isDecaySlot, float const* staticTopP, + int64_t const* slots, int32_t numRows, cudaStream_t stream) +{ + if (numRows == 0) + { + return; + } + constexpr int32_t kBlock = 256; + int32_t const grid = (numRows + kBlock - 1) / kBlock; + toppDecayGatherKernel<<>>(rowTopP, runtimeTopP, isDecaySlot, staticTopP, slots, numRows); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h new file mode 100644 index 000000000000..5e3eb2ead375 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "tensorrt_llm/common/config.h" +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! \brief Fused post-sample runtime top-p update for the PyTorch TorchSampler's +//! Top-P Decay feature. One thread per sampled row; the decay-active gate is +//! applied on-device via \p isDecaySlot so no host-side filtering is needed. +//! +//! Unlike the legacy invokeComputeToppDecay (samplingTopPKernels.h), this variant +//! gathers the sampled token in-kernel from a slot-indexed strided view of the +//! new-tokens buffer (\p stepTokens with element stride \p stepTokenStride, i.e. +//! new_tokens[step, slot, beam] for a fixed step/beam), applies a gated reset +//! (reset_id < 0 disables the reset), and filters decay-active slots on-device +//! via \p isDecaySlot. +//! +//! For each sampled row i with slot s = sampledSlots[i] where isDecaySlot[s]: +//! runtimeTopP[s] = (resetIds[s] >= 0 && stepTokens[s * stride] == resetIds[s]) +//! ? initialTopP[s] +//! : max(runtimeTopP[s] * topPDecay[s], topPMin[s]) +//! +//! All per-slot arrays are length numSlots; \p sampledSlots is length numSampled. +//! \p runtimeTopP is updated in place. +void invokeToppDecayUpdate(float* runtimeTopP, float const* initialTopP, float const* topPDecay, float const* topPMin, + int32_t const* resetIds, bool const* isDecaySlot, int32_t const* stepTokens, int64_t stepTokenStride, + int64_t const* sampledSlots, int32_t numSampled, cudaStream_t stream); + +//! \brief Fused pre-sample per-row top-p gather for the Top-P Decay feature. +//! Replaces the eager index_select(runtime) + index_select(gate) + where(static) +//! chain with a single launch: +//! rowTopP[i] = isDecaySlot[slots[i]] ? runtimeTopP[slots[i]] : staticTopP[i] +//! \p rowTopP / \p staticTopP / \p slots are length numRows (per-step group rows); +//! \p runtimeTopP / \p isDecaySlot are per-slot arrays. +void invokeToppDecayGather(float* rowTopP, float const* runtimeTopP, bool const* isDecaySlot, float const* staticTopP, + int64_t const* slots, int32_t numRows, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index b95426f580fe..61912effeba2 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -113,6 +113,7 @@ add_library( IndexerKCacheGatherOp.cpp IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp + samplerOp.cpp mlaRopeInplaceOp.cpp ncclCommunicatorOp.cpp allocateOutput.cpp diff --git a/cpp/tensorrt_llm/thop/samplerOp.cpp b/cpp/tensorrt_llm/thop/samplerOp.cpp new file mode 100644 index 000000000000..ee39b5c44df8 --- /dev/null +++ b/cpp/tensorrt_llm/thop/samplerOp.cpp @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Torch custom ops for the PyTorch backend sampler (TorchSampler). Aggregates +// sampler-related ops here so follow-up work can add new sampler ops alongside +// existing ones, mirroring how specDecOp.cpp groups the speculative-decoding ops. + +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +namespace th = torch; +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//////////////////////////////////////////////////////////////////////////////////////////////////////////// +// Fused post-sample runtime top-p update for the TorchSampler Top-P Decay +// feature. Applies, in place, for every sampled row whose slot is decay-active: +// runtime_top_p[slot] = (reset_id >= 0 && last_token == reset_id) +// ? initial_top_p[slot] +// : max(runtime_top_p[slot] * top_p_decay[slot], top_p_min[slot]) +// The decay-active gate is applied on-device via is_decay_slot, so the hot path +// needs no host-side .tolist() / set intersection. +void top_p_decay_update(th::Tensor runtime_top_p, th::Tensor initial_top_p, th::Tensor top_p_decay, + th::Tensor top_p_min, th::Tensor reset_ids, th::Tensor is_decay_slot, th::Tensor step_tokens, + th::Tensor sampled_slots) +{ + TORCH_CHECK(runtime_top_p.is_cuda() && initial_top_p.is_cuda() && top_p_decay.is_cuda() && top_p_min.is_cuda() + && reset_ids.is_cuda() && is_decay_slot.is_cuda() && step_tokens.is_cuda() && sampled_slots.is_cuda(), + "all top_p_decay_update tensors must be CUDA tensors"); + + TORCH_CHECK(runtime_top_p.scalar_type() == th::kFloat32, "runtime_top_p must be float32"); + TORCH_CHECK(initial_top_p.scalar_type() == th::kFloat32, "initial_top_p must be float32"); + TORCH_CHECK(top_p_decay.scalar_type() == th::kFloat32, "top_p_decay must be float32"); + TORCH_CHECK(top_p_min.scalar_type() == th::kFloat32, "top_p_min must be float32"); + TORCH_CHECK(reset_ids.scalar_type() == th::kInt32, "reset_ids must be int32"); + TORCH_CHECK(is_decay_slot.scalar_type() == th::kBool, "is_decay_slot must be bool"); + TORCH_CHECK(step_tokens.scalar_type() == th::kInt32, "step_tokens must be int32"); + TORCH_CHECK(sampled_slots.scalar_type() == th::kInt64, "sampled_slots must be int64"); + + // step_tokens is a slot-indexed 1-D view of the new-tokens buffer for a fixed + // step/beam (new_tokens[step, :, beam]); it is strided, not contiguous. The + // kernel gathers tokens through its element stride, avoiding a separate + // gather + cast launch on the hot path. + TORCH_CHECK(runtime_top_p.is_contiguous() && initial_top_p.is_contiguous() && top_p_decay.is_contiguous() + && top_p_min.is_contiguous() && reset_ids.is_contiguous() && is_decay_slot.is_contiguous() + && sampled_slots.is_contiguous(), + "all top_p_decay_update tensors (except step_tokens) must be contiguous"); + TORCH_CHECK(step_tokens.dim() == 1, "step_tokens must be a 1-D (strided) view"); + + auto const numSlots = runtime_top_p.size(0); + TORCH_CHECK(initial_top_p.size(0) == numSlots && top_p_decay.size(0) == numSlots && top_p_min.size(0) == numSlots + && reset_ids.size(0) == numSlots && is_decay_slot.size(0) == numSlots, + "all per-slot tensors must have the same length (max_num_sequences)"); + TORCH_CHECK(step_tokens.size(0) >= numSlots, "step_tokens must cover all slots"); + + auto const numSampled = sampled_slots.size(0); + + auto stream = at::cuda::getCurrentCUDAStream(runtime_top_p.get_device()); + tk::invokeToppDecayUpdate(runtime_top_p.data_ptr(), initial_top_p.data_ptr(), + top_p_decay.data_ptr(), top_p_min.data_ptr(), reset_ids.data_ptr(), + is_decay_slot.data_ptr(), step_tokens.data_ptr(), step_tokens.stride(0), + sampled_slots.data_ptr(), static_cast(numSampled), stream); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////////////// +// Fused pre-sample per-row top-p gather for the Top-P Decay feature. Returns +// row_top_p[i] = is_decay_slot[slots[i]] ? runtime_top_p[slots[i]] : static_top_p[i] +// replacing the index_select(runtime) + index_select(gate) + where(static) chain +// with a single launch. +th::Tensor top_p_decay_gather( + th::Tensor runtime_top_p, th::Tensor is_decay_slot, th::Tensor static_top_p, th::Tensor slots) +{ + TORCH_CHECK(runtime_top_p.is_cuda() && is_decay_slot.is_cuda() && static_top_p.is_cuda() && slots.is_cuda(), + "all top_p_decay_gather tensors must be CUDA tensors"); + TORCH_CHECK(runtime_top_p.scalar_type() == th::kFloat32, "runtime_top_p must be float32"); + TORCH_CHECK(is_decay_slot.scalar_type() == th::kBool, "is_decay_slot must be bool"); + TORCH_CHECK(static_top_p.scalar_type() == th::kFloat32, "static_top_p must be float32"); + TORCH_CHECK(slots.scalar_type() == th::kInt64, "slots must be int64"); + TORCH_CHECK(runtime_top_p.is_contiguous() && is_decay_slot.is_contiguous() && static_top_p.is_contiguous() + && slots.is_contiguous(), + "all top_p_decay_gather tensors must be contiguous"); + TORCH_CHECK(runtime_top_p.size(0) == is_decay_slot.size(0), "per-slot tensors must have the same length"); + auto const numRows = slots.size(0); + TORCH_CHECK(static_top_p.size(0) == numRows, "static_top_p and slots must have the same length"); + // NB: every slots[i] must lie in [0, runtime_top_p.size(0)); the kernel indexes the per-slot + // arrays with it unchecked (values live on device, so validating here would force a sync). + // The caller (TorchSampler) guarantees this: slots come from seq_slots, which is bounded by + // max_num_sequences -- the allocation size of the per-slot store tensors. + + auto row_top_p = th::empty_like(static_top_p); + auto stream = at::cuda::getCurrentCUDAStream(runtime_top_p.get_device()); + tk::invokeToppDecayGather(row_top_p.data_ptr(), runtime_top_p.data_ptr(), + is_decay_slot.data_ptr(), static_top_p.data_ptr(), slots.data_ptr(), + static_cast(numRows), stream); + return row_top_p; +} + +} // end namespace torch_ext + +TRTLLM_NAMESPACE_END + +//////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "top_p_decay_update(Tensor(a!) runtime_top_p, Tensor initial_top_p, Tensor top_p_decay, Tensor top_p_min, " + "Tensor reset_ids, Tensor is_decay_slot, Tensor step_tokens, Tensor sampled_slots) -> ()"); + m.def( + "top_p_decay_gather(Tensor runtime_top_p, Tensor is_decay_slot, Tensor static_top_p, Tensor slots) " + "-> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("top_p_decay_update", &tensorrt_llm::torch_ext::top_p_decay_update); + m.impl("top_p_decay_gather", &tensorrt_llm::torch_ext::top_p_decay_gather); +} diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index d5a9eaef0685..437b68aaf1bf 100644 --- a/docs/source/features/sampling.md +++ b/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,18 @@ 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). + + * 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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py new file mode 100644 index 000000000000..70ece23a75d3 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Thin Python wrappers over the trtllm C++ sampling ops used by TorchSampler. + +These forward to ``torch.ops.trtllm.*`` custom ops (registered from +``cpp/tensorrt_llm/thop``); the wrappers exist to give the sampler a stable, +keyword-only Python surface decoupled from the raw op-registration names. +""" + +import torch + + +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. + + For each sampled row whose slot is decay-active (per ``is_decay_slot``):: + + runtime_top_p[slot] = initial_top_p[slot] if reset_id >= 0 and token == reset_id + = max(runtime_top_p[slot] * decay, top_p_min) otherwise + + 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]``) -- + the kernel gathers each sampled token in-kernel; ``sampled_slots`` is 1-D of + length ``num_sampled`` (this iteration's rows). ``runtime_top_p`` is mutated + in place; nothing is returned. + """ + torch.ops.trtllm.top_p_decay_update( + runtime_top_p, + initial_top_p, + top_p_decay, + top_p_min, + reset_ids, + is_decay_slot, + step_tokens, + sampled_slots, + ) + + +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 + + replacing the eager ``index_select`` x2 + ``where`` chain with one launch. + """ + return torch.ops.trtllm.top_p_decay_gather(runtime_top_p, is_decay_slot, static_top_p, slots) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index bb50aab8c351..1b28a6bb8338 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -88,6 +88,7 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests +from .ops.trtllm import top_p_decay_update from .sampling_utils import ( BEAM_SEARCH_PAD_TOKEN, GREEDY, @@ -97,11 +98,13 @@ GenericStrategyKeyType, Strategy, StrategyMetadata, + TopPDecayOverride, UtilsSamplingParams, get_rejected_indices, resolve_sampling_strategy, sample, sample_rejected, + top_p_decay_active, ) if sys.version_info[:2] >= (3, 12): @@ -478,6 +481,14 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) + # These top-p-decay fields live on the C++ SamplingConfig as optional> + # (a shape designed for the batched TRT-LLM sampler); the torch sampler consumes + # them per request, so we unwrap the singleton list into a scalar here. When the + # TRT-LLM sampler is removed, this SamplingConfig-based plumbing should be removed + # too in favor of reading the values directly from the per-request params. + top_p_decay = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_decay)) + top_p_min = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_min)) + top_p_reset_ids = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_p_reset_ids)) beam_width_out = _get_beam_width_out(request) beam_width_in = _get_beam_width_in(request) use_beam_search = _get_max_beam_width(request) > 1 @@ -489,6 +500,9 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: beam_width_in=beam_width_in, beam_width_out=beam_width_out, use_beam_search=use_beam_search, + top_p_decay=top_p_decay, + top_p_min=top_p_min, + top_p_reset_ids=top_p_reset_ids, ) @@ -788,6 +802,10 @@ def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: return torch.empty(shape, dtype=torch.int, device=device) +def float_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: + return torch.empty(shape, dtype=torch.float32, device=device) + + @dataclass(kw_only=True, frozen=True) class _BatchedSamplingResult: # Original request indices for all requests (permuted due to batching by strategy): @@ -2247,6 +2265,31 @@ class LogProbsStore: """Shape: batch_size, max_tokens, max_topk_logprobs Usage: Stores the values of the topk logprobs""" + @dataclass(kw_only=True) + class TopPDecayStore: + """Per-slot runtime state for Top-P Decay (FlashInfer path). + + Each tensor has shape (max_num_sequences,). Only slots that are members + of ``TorchSampler._top_p_decay_slots`` hold meaningful values; all reads + and updates are gated by that set, so stale entries from a prior occupant + of a reused slot are never consumed. + """ + + runtime_top_p_decay_cuda: torch.Tensor + """The current (decaying) top-p per slot; mutated post-sample each step.""" + initial_top_p_decay_cuda: torch.Tensor + """The initial top-p per slot; used to reset on a reset-id match.""" + top_p_decay_cuda: torch.Tensor + """Per-slot multiplicative decay factor.""" + top_p_decay_min_cuda: torch.Tensor + """Per-slot lower bound for the decayed top-p.""" + top_p_decay_reset_ids_cuda: torch.Tensor + """Per-slot reset token id (< 0 never matches a sampled token).""" + is_top_p_decay_slot_cuda: torch.Tensor + """Per-slot bool gate (device mirror of ``_top_p_decay_slots``). Lets the + fused post-sample update kernel filter decay-active slots on the GPU, + avoiding a host-side ``.tolist()`` / set intersection each step.""" + @dataclass(kw_only=True) class Store: new_tokens: torch.Tensor @@ -2258,6 +2301,8 @@ class Store: """Holds data related to beam search.""" log_probs_store: "TorchSampler.LogProbsStore" """Holds data related to log-probs handling.""" + top_p_decay_store: "TorchSampler.TopPDecayStore" + """Holds per-slot runtime state for Top-P Decay.""" def _create_store(self) -> Store: # Tensors necessary for all sampling methods @@ -2308,10 +2353,25 @@ def _create_store(self) -> Store: seq_offsets=seq_offsets, beam_idx_arange=beam_idx_arange, ) + # Per-slot Top-P Decay runtime state (FlashInfer path). Allocated for all + # sampler instances; only slots in self._top_p_decay_slots are ever read. + n = (self.max_num_sequences,) + top_p_decay_store = self.TopPDecayStore( + runtime_top_p_decay_cuda=float_tensor(n), + initial_top_p_decay_cuda=float_tensor(n), + top_p_decay_cuda=float_tensor(n), + top_p_decay_min_cuda=float_tensor(n), + top_p_decay_reset_ids_cuda=int_tensor(n), + # The gate buffer IS the gate, so (unlike the others) it must start + # False: a slot that was never admitted as decay-active must read False. + is_top_p_decay_slot_cuda=torch.zeros(n, dtype=torch.bool, device="cuda"), + ) + return self.Store( new_tokens=new_tokens, log_probs_store=log_probs_store, beam_search_store=beam_search_store, + top_p_decay_store=top_p_decay_store, ) @dataclass(frozen=True, kw_only=True) @@ -2361,6 +2421,10 @@ def __init__(self, args: Args): "in requirements.txt." ) self._grouped_sampler_cls = FlashInferGroupedStrategySampler + # Slots with an active top-p-decay request. Sole gate for reading/updating + # the per-slot TopPDecayStore buffers; discarded on slot reuse so stale + # buffer entries are never consumed. + self._top_p_decay_slots: set[int] = set() # AutoDeploy build creates the sampler in inference mode, # which would disallow in-place mutating of new_tokens. @@ -2461,6 +2525,42 @@ def get_spec_tree_manager( def _use_beam_search(self) -> bool: return self.max_beam_width > 1 + def _validate_top_p_decay_request(self, request: LlmRequest) -> None: + """Reject unsupported combinations for a top-p-decay-active request. + + Top-p decay is supported only for single-token decode steps without beam + search. Called from validate_request (request admission), so a violating + request is failed individually instead of aborting the whole batch. + """ + params = _request_get_sampling_params(request) + if not top_p_decay_active(params): + return + if params.use_beam_search: + raise ValueError("top_p_decay is not supported with beam search.") + # A non-zero draft length means the request carries speculative draft + # tokens and produces multiple tokens per step (req_num_steps = + # 1 + draft_token_length). One-model speculation (vanilla MTP, one-model + # Eagle3 / MTP-Eagle, SA, draft-target-one-model) uses its own + # SpecSamplerBase-derived sampler and never reaches TorchSampler; the + # drafter-based modes that DO flow draft tokens through TorchSampler + # (two-model draft-target, NGram, user-provided, two-model Eagle3 / + # MTP-Eagle) are what can make this length non-zero. top-p decay does not + # support these multi-token steps. + # NB: at admission time the draft tokens of drafter-based modes are + # usually not attached yet, so this check is best-effort. Two-model + # speculation (the only source of such requests in TorchSampler) is + # slated for removal, so in practice no speculative request reaches the + # decay path; a debug assert in _build_top_p_runtime_override guards the + # invariant at sample time. + if get_draft_token_length(request) > 0: + raise ValueError( + "top_p_decay is not supported for requests carrying speculative " + "draft tokens (req_num_steps > 1). This covers the drafter-based " + "modes routed through TorchSampler (two-model draft-target, NGram, " + "user-provided, two-model Eagle3 / MTP-Eagle); one-model " + "speculation uses its own sampler and is unaffected." + ) + def _can_use_fast_greedy_path(self, requests: list[LlmRequest]) -> bool: """ Check if we can use the fast argmax path for greedy sampling. @@ -2921,6 +3021,10 @@ def _collect_new_requests_for_setup( @override def validate_request(self, request: LlmRequest) -> None: + # Reject unsupported top-p-decay combinations at admission, so only the + # offending request fails (raising later, inside setup_sampler_step or + # sampling, would abort the whole executor step). + self._validate_top_p_decay_request(request) if self._use_beam_search: if request.py_return_log_probs: if request.py_num_logprobs > 1: @@ -2949,6 +3053,12 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: seq_slots: list[int] = [] # Used for beam search updates max_prompt_len: int = 0 + # Used for top-p-decay store updates (only populated for decay-active reqs) + decay_seq_slots: list[int] = [] + decay_initial_top_p: list[float] = [] + decay_top_p_decay: list[float] = [] + decay_top_p_min: list[float] = [] + decay_top_p_reset_ids: list[int] = [] # Prepare finish reasons handler self._finish_reasons_handler.setup_new_request_handling() @@ -2956,6 +3066,29 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: slot = request.py_seq_slot assert slot is not None seq_slots.append(slot) + # Drop any stale top-p-decay membership from a prior occupant of this + # slot; re-added below only if this request itself uses decay. + # (Unsupported decay combinations were already rejected per-request in + # validate_request at admission.) + self._top_p_decay_slots.discard(slot) + decay_params = _request_get_sampling_params(request) + if top_p_decay_active(decay_params): + self._top_p_decay_slots.add(slot) + decay_seq_slots.append(slot) + # Initial runtime top-p defaults to 1.0 when top_p is unset. + decay_initial_top_p.append( + decay_params.top_p if decay_params.top_p is not None else 1.0 + ) + # decay is guaranteed non-None and < 1.0 here (top_p_decay_active); + # min/reset fall back to the C++ runtime defaults when unset. + assert decay_params.top_p_decay is not None + decay_top_p_decay.append(decay_params.top_p_decay) + decay_top_p_min.append( + decay_params.top_p_min if decay_params.top_p_min is not None else 1e-6 + ) + decay_top_p_reset_ids.append( + decay_params.top_p_reset_ids if decay_params.top_p_reset_ids is not None else -1 + ) # update temp_data with this requests data self._finish_reasons_handler.prepare_for_new_request(request) @@ -2995,6 +3128,21 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: all_sampling_requests=new_requests + scheduled_requests.generation_requests, ) + # Clear the device decay gate for every newly-admitted slot (covers slot + # reuse: a slot previously decay-active but reused by a non-decay request + # must read False). Decay-active slots are then set True below. + decay_gate = self.store.top_p_decay_store.is_top_p_decay_slot_cuda + decay_gate.index_fill_(0, seq_slots_tensor_cuda_long, False) + + if decay_seq_slots: + self._update_top_p_decay_store_for_new_requests( + decay_seq_slots=decay_seq_slots, + initial_top_p=decay_initial_top_p, + top_p_decay=decay_top_p_decay, + top_p_min=decay_top_p_min, + top_p_reset_ids=decay_top_p_reset_ids, + ) + if self._use_beam_search: beam_search_store = self.store.beam_search_store assert beam_search_store is not None @@ -3005,6 +3153,46 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: max_prompt_len=max_prompt_len, ) + def _update_top_p_decay_store_for_new_requests( + self, + *, + decay_seq_slots: list[int], + initial_top_p: list[float], + top_p_decay: list[float], + top_p_min: list[float], + top_p_reset_ids: list[int], + ) -> None: + """Initialize per-slot Top-P Decay buffers for newly admitted decay requests. + + runtime_top_p and initial_top_p both start at the effective initial top-p; + the runtime value is decayed post-sample each step. + """ + store = self.store.top_p_decay_store + device = store.runtime_top_p_decay_cuda.device + slots_cuda = torch.tensor( + decay_seq_slots, device="cpu", dtype=torch.int64, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + floats_host = torch.tensor( + [initial_top_p, top_p_decay, top_p_min], + device="cpu", + dtype=torch.float32, + pin_memory=prefer_pinned(), + ) + floats_cuda = floats_host.to(device, non_blocking=True) + initial_cuda = floats_cuda[0] + reset_ids_cuda = torch.tensor( + top_p_reset_ids, device="cpu", dtype=torch.int32, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + + store.runtime_top_p_decay_cuda.index_copy_(0, slots_cuda, initial_cuda) + store.initial_top_p_decay_cuda.index_copy_(0, slots_cuda, initial_cuda) + store.top_p_decay_cuda.index_copy_(0, slots_cuda, floats_cuda[1]) + store.top_p_decay_min_cuda.index_copy_(0, slots_cuda, floats_cuda[2]) + store.top_p_decay_reset_ids_cuda.index_copy_(0, slots_cuda, reset_ids_cuda) + # Enable the device gate for these decay-active slots (cleared for all new + # slots in setup_sampler_step just before this call). + store.is_top_p_decay_slot_cuda.index_fill_(0, slots_cuda, True) + @staticmethod def _prepare_beam_search( beam_search_store: BeamSearchStore, @@ -3682,6 +3870,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: for req_idx, req in enumerate(state.requests): if req.state == LlmRequestState.GENERATION_COMPLETE: + self._drop_top_p_decay_slot_if_finished(req) continue if req.py_beam_width > 1: @@ -3748,6 +3937,25 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: self._finish_reasons_handler.store.num_accepted_draft_tokens_host[ req.py_seq_slot ] = req.py_num_accepted_draft_tokens + self._drop_top_p_decay_slot_if_finished(req) + + def _drop_top_p_decay_slot_if_finished(self, req: LlmRequest) -> None: + """Retire a finished request's slot from the top-p-decay membership set. + + Without this, `_top_p_decay_slots` would only shrink on slot reuse, so the + O(1) early-outs on the sampling hot path (post-sample update, override + build) would stay disabled forever once any decay request had run. The + device-side buffers need no cleanup: a freed slot is never sampled, and + slot reuse re-initializes them in setup_sampler_step. Requests that finish + outside the sampler (e.g. cancellation) are covered by the slot-reuse + cleanup instead. + """ + if ( + self._top_p_decay_slots + and req.state == LlmRequestState.GENERATION_COMPLETE + and req.py_seq_slot is not None + ): + self._top_p_decay_slots.discard(req.py_seq_slot) def _return_log_probs(self, requests: list[LlmRequest]) -> bool: return any(req.py_return_log_probs for req in requests) @@ -4015,6 +4223,79 @@ def provision_bias_index() -> int: # sharing). logits[logits_bias_mask_cuda] += biases_tensor_cuda + def _build_top_p_runtime_override( + self, + *, + strategy_key: Any, + group_req_indices: torch.Tensor, + req_num_steps: torch.Tensor, + seq_slots: torch.Tensor, + seq_slots_cuda: torch.Tensor, + device: torch.device, + ) -> Optional[TopPDecayOverride]: + """Build the per-row runtime top-p override for a sampling group. + + Returns None unless this is a top-p / top_k_top_p group that may contain + a top-p-decay-active row. The override's ``slots`` tensor is aligned to + the group's per-STEP row order (matching group_strategies_per_step); + decay-active rows are single-token, while non-decay rows (possibly + multi-step draft rows) are gated out on-device by ``is_decay_slot``. + """ + if strategy_key not in ("top_p", "top_k_top_p"): + return None + if not self._top_p_decay_slots: + return None + store = self.store.top_p_decay_store + # Fast path (steady-state decoding): if every row in the group is + # single-token, the per-STEP row order equals the per-request order, and + # (group_req_indices being sorted ascending) a contiguous group's slots + # are just a slice of seq_slots_cuda -- no host layout build and no H2D + # copy. Decay presence in the group is not checked host-side here: a + # group without decay rows is gated out on-device by is_decay_slot, so + # every row samples with its static top-p -- same result as returning + # None, minus the short-circuit. + first_req = int(group_req_indices[0]) + last_req = int(group_req_indices[-1]) + if last_req - first_req + 1 == group_req_indices.size(0) and ( + int(req_num_steps[group_req_indices].max()) == 1 + ): + per_step_slots_cuda = seq_slots_cuda[first_req : last_req + 1] + else: + group_seq_slots = seq_slots[group_req_indices].tolist() + group_steps = req_num_steps[group_req_indices].tolist() + # Build a per-STEP slot layout aligned with group_strategies_per_step + # (each request contributes req_num_steps rows). The per-row decay mask is + # not built here: it is gathered on-device from the is_decay_slot gate + # below. A decay-active row is always single-token: the only source of + # multi-step rows in TorchSampler is two-model speculation (slated for + # removal), and decay + draft tokens is rejected per-request at + # admission in validate_request -- so this is an internal invariant, + # not a user-reachable error path. + per_step_slots: list[int] = [] + has_decay = False + for slot, steps in zip(group_seq_slots, group_steps): + if slot in self._top_p_decay_slots: + has_decay = True + assert steps == 1, ( + "top_p_decay row with req_num_steps != 1; decay + draft tokens " + "should have been rejected at admission" + ) + per_step_slots.extend([slot] * steps) + if not has_decay: + return None + per_step_slots_cuda = torch.tensor( + per_step_slots, device="cpu", dtype=torch.int64, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + # The actual per-row gather (runtime top-p gated by is_decay_slot, the + # device mirror of _top_p_decay_slots kept in sync by setup_sampler_step's + # clear-then-set) happens in one fused launch inside + # _apply_top_p_runtime_override, where the static per-row top-p lives. + return TopPDecayOverride( + slots=per_step_slots_cuda, + runtime_top_p=store.runtime_top_p_decay_cuda, + is_decay_slot=store.is_top_p_decay_slot_cuda, + ) + @nvtx_range("sample_batched_by_strategy") @torch.inference_mode() def _sample_batched_by_strategy( @@ -4173,6 +4454,24 @@ def _sample_batched_by_strategy( for _ in range(steps) ] + # For top-p / top_k_top_p groups, source the runtime (decayed) top-p + # for any decay-active rows from the per-slot store. Decay is single-token + # only, so per-step order matches per-request order here. Returns None + # on the non-FlashInfer path or when no decay row can be present, in + # which case the top-p-decay kwarg is omitted entirely (the Simple + # sampler backend does not accept it). + decay_override = self._build_top_p_runtime_override( + strategy_key=strategy_key, + group_req_indices=group_req_indices, + req_num_steps=req_num_steps, + seq_slots=seq_slots, + seq_slots_cuda=seq_slots_cuda, + device=group_logits_cuda.device, + ) + top_p_decay_kwargs = ( + {} if decay_override is None else {"decay_override": decay_override} + ) + group_next_tokens_cuda, group_softmax_cuda, group_temperature_cuda = ( self._grouped_sampler_cls.sample_grouped_strategies( strategy_key, @@ -4182,6 +4481,7 @@ def _sample_batched_by_strategy( return_probs=needs_probs, group_logit_indices=logit_indices_for_sampler, group_metadata=group_metadata, + **top_p_decay_kwargs, ) ) batch_next_tokens_offset_end = ( @@ -4284,6 +4584,7 @@ def _unbatch_sampling_results( new_tokens_cuda: torch.Tensor, req_num_generated_tokens: torch.Tensor, seq_slots: torch.Tensor, + seq_slots_cuda: torch.Tensor, ) -> torch.Tensor: batch_req_indices = batched_sampling_result.batch_req_indices batch_next_tokens_cuda_int = batched_sampling_result.batch_next_tokens_cuda_int @@ -4317,8 +4618,58 @@ def _dims_canonically_ordered(t: torch.Tensor) -> bool: new_tokens_cuda.view(-1, *new_tokens_cuda.shape[2:]).scatter_( 0, batch_dest_indices_1d_cuda, batch_next_tokens_cuda_int ) + # Post-sample: decay the runtime top-p for any decay-active slots that were + # sampled this iteration (must run after tokens land in new_tokens_cuda). + # batch_req_indices is a permutation of all sampled requests, so the set of + # sampled slots is exactly seq_slots (the kernel updates each slot + # independently; order is irrelevant) -- pass the resident device copy + # instead of gathering seq_slots[batch_req_indices] on host and copying it. + self._update_top_p_decay_after_sample( + new_tokens_cuda=new_tokens_cuda, + sampled_slots_cuda=seq_slots_cuda, + ) return self._copy_to_host(new_tokens_cuda) + def _update_top_p_decay_after_sample( + self, + *, + new_tokens_cuda: torch.Tensor, + sampled_slots_cuda: torch.Tensor, + ) -> None: + """Apply the C++ ``computeToppDecay`` update for sampled decay-active slots. + + For each slot sampled this iteration that is decay-active: + runtime_p = initial_p if sampled_token == reset_id and reset_id >= 0 + runtime_p = max(runtime_p * decay, top_p_min) otherwise + + Restricting to the sampled slots avoids reading stale new_tokens_cuda + for slots that were not scheduled this iteration. + """ + # Host-side O(1) early-out: skip entirely when no request uses decay + # (avoids launching the kernel). The per-slot decay-active filtering is + # otherwise done on-device via the is_decay_slot gate, so no .tolist() / + # set intersection is needed on the hot path. + if not self._top_p_decay_slots: + return + store = self.store.top_p_decay_store + # Single fused kernel: gates on is_decay_slot on-device, gathers the + # sampled token in-kernel from the slot-indexed strided view of this + # iteration's new-tokens buffer (decay is single-token-only, so the + # token is at local step 0, beam 0), and applies + # runtime_p = (reset_id>=0 && tok==reset_id) ? initial_p + # : max(runtime_p*decay, top_p_min) + # in place -- no separate gather/cast launches on the hot path. + top_p_decay_update( + runtime_top_p=store.runtime_top_p_decay_cuda, + initial_top_p=store.initial_top_p_decay_cuda, + top_p_decay=store.top_p_decay_cuda, + top_p_min=store.top_p_decay_min_cuda, + reset_ids=store.top_p_decay_reset_ids_cuda, + is_decay_slot=store.is_top_p_decay_slot_cuda, + step_tokens=new_tokens_cuda[DEFAULT_STEP_IDX, :, DEFAULT_BEAM_IDX], + sampled_slots=sampled_slots_cuda, + ) + @staticmethod @torch.inference_mode() def _apply_min_length_penalty( @@ -4953,6 +5304,7 @@ def _process_requests( new_tokens_cuda=new_tokens_cuda, req_num_generated_tokens=sampling_requests_metadata.req_num_generated_tokens, seq_slots=seq_slots_host, + seq_slots_cuda=seq_slots_cuda, ) # NB: update_requests syncs w/ device computation and async D2H copies diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index f99946ffd5a1..6a7e6c2ef478 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -39,6 +39,7 @@ top_p_renorm_probs_op, top_p_sampling_from_probs_op, ) +from tensorrt_llm._torch.pyexecutor.sampler.ops.trtllm import top_p_decay_gather from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import ( GREEDY_TEMPERATURE_THRESHOLD, BeamSearchMetadata, @@ -106,6 +107,10 @@ class UtilsSamplingParams: use_beam_search: Whether to use beam search. beam_width_in: The beam_width of a request before the sampling step. beam_width_out: The beam_width of a request after the sampling step. + top_p_decay: Per-step multiplicative decay applied to the runtime top-p. + top_p_min: Lower bound for the decayed runtime top-p. + top_p_reset_ids: Token id which, when sampled, resets the runtime top-p to + its initial value. A value < 0 never matches a token. """ temperature: Optional[float] @@ -114,6 +119,36 @@ class UtilsSamplingParams: use_beam_search: Optional[bool] beam_width_in: Optional[int] = None beam_width_out: Optional[int] = None + top_p_decay: Optional[float] = None + top_p_min: Optional[float] = None + top_p_reset_ids: Optional[int] = None + + +@dataclass(frozen=True, kw_only=True) +class TopPDecayOverride: + """Per-group runtime top-p override for decay-active rows. + + ``slots`` maps each per-step group row to its sequence slot; the decayed + per-row top-p is gathered on-device from the per-slot ``runtime_top_p`` + store, gated by ``is_decay_slot`` (non-decay rows keep their static top-p). + """ + + slots: torch.Tensor + """Per-step group rows' sequence slots (int64, device).""" + runtime_top_p: torch.Tensor + """Per-slot runtime (decayed) top-p store (float32, device).""" + is_decay_slot: torch.Tensor + """Per-slot decay-active gate (bool, device).""" + + +def top_p_decay_active(params: UtilsSamplingParams) -> bool: + """Whether dynamic top-p decay is active for a request. + + Decay is active iff ``top_p_decay`` is explicitly set and ``< 1.0``. A decay + of ``1.0`` (the C++ default) is a no-op, and ``top_p_min`` / ``top_p_reset_ids`` + alone do not activate dynamic behavior. + """ + return params.top_p_decay is not None and params.top_p_decay < 1.0 def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) -> Strategy: @@ -124,13 +159,25 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - top_p = params.top_p top_k = params.top_k + # When top-p decay is active, the runtime top-p can shrink below 1.0 on later + # steps even if the request starts at top_p == 1.0 / None. Force a top-p-capable + # strategy (defaulting the initial top-p to 1.0) so the decayed value takes effect. + # However, an EXPLICIT greedy control (see params_imply_explicit_greedy) is a + # deterministic intent that must be honored even with decay set -- only the + # implicit "all params unset" greedy case is overridden by decay. + decay_active = top_p_decay_active(params) + if SamplingParams.params_imply_greedy_decoding( temperature=temperature, top_p=top_p, top_k=top_k, use_beam_search=use_beam_search, ): - return GREEDY + explicitly_greedy = SamplingParams.params_imply_explicit_greedy( + temperature=temperature, top_p=top_p, top_k=top_k + ) + if explicitly_greedy or not decay_active: + return GREEDY # --- resolving default values # NB: not greedy, hence temperature != 0 if specified @@ -152,7 +199,10 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - assert top_k > 1, "non-greedy sampling requires valid top_k" need_top_k = top_k < vocab_size assert top_p > 0, "non-greedy sampling requires valid top_p" - need_top_p = top_p < 1 + # A decay-active request must go through a top-p-capable path even when its + # initial top_p is 1.0, so the runtime top-p (sourced per-row at sample time) + # can shrink the nucleus on later steps. + need_top_p = top_p < 1 or decay_active if need_top_p: if need_top_k: @@ -273,6 +323,40 @@ def _make_tensor(data: list[Any], dtype: torch.dtype, device: torch.device) -> t device=device, non_blocking=True ) + def _apply_top_p_runtime_override( + self, + decay_override: "TopPDecayOverride", + ) -> None: + """Override the per-row static top-p with the decayed top-p. + + Only decay-active rows (per the on-device ``is_decay_slot`` gate) are + overridden, so a group mixing top-p-decay and plain top-p requests + keeps each row's correct value. The overridden ``self._top_p`` tensor + then feeds both sampling and ``top_p_renorm_probs_op`` (so processed + logprobs match). Single fused launch (gather + gate + select). + + The caller (sample_grouped_strategies) is responsible for the + None-check; this method requires an actual override. + """ + # Only TopP*/TopKTopP* impls (which own a per-row ``_top_p`` tensor) + # receive a runtime override; access via getattr keeps this base-class + # helper independent of subclass attribute declarations. + static_top_p: torch.Tensor = getattr(self, "_top_p") + assert static_top_p.shape == decay_override.slots.shape, ( + static_top_p.shape, + decay_override.slots.shape, + ) + setattr( + self, + "_top_p", + top_p_decay_gather( + runtime_top_p=decay_override.runtime_top_p, + is_decay_slot=decay_override.is_decay_slot, + static_top_p=static_top_p, + slots=decay_override.slots, + ), + ) + @staticmethod def _prepare_logits_with_temperature( logits: torch.Tensor, @@ -768,6 +852,7 @@ def sample_grouped_strategies( generator: Optional[torch.Generator] = None, return_probs: bool, group_metadata: StrategyMetadata | None = None, + decay_override: Optional[TopPDecayOverride] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: """Sample grouped strategies. @@ -817,6 +902,10 @@ def sample_grouped_strategies( else: assert group_logit_indices.size(0) == beam_width_in * len(strategies) strategy_impl = strategy_impl_cls.from_strategies(strategies, cuda_device=logits.device) + # For top-p / top_k_top_p groups, override the per-row static top-p with the + # decayed top-p on the decay-active rows (gated on-device). + if decay_override is not None: + strategy_impl._apply_top_p_runtime_override(decay_override) next_tokens, softmax = strategy_impl.sample( logits, group_logit_indices=group_logit_indices, diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index bded9477cd3d..b19576964c16 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -211,7 +211,7 @@ class SamplingParams: If temperature > 0 and/or top_k > 1 are specified, sampling will proceed accordingly and top_p will default to top_p = 1. Setting top_p = 0 should result in greedy sampling, but is currently disallowed in the backend. top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. None means using C++ runtime default 1.e-6. Defaults to None. - top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. Indicates where to reset the decay. None means using C++ runtime default 1. Defaults to None. + top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. The token id which, when sampled, resets the decayed top-P to its initial value. None means using C++ runtime default -1 (which never matches a token). Defaults to None. top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. None means using C++ runtime default 1.f. Defaults to None. seed (int, optional): Controls the random seed used by the random number generator in sampling. None means using C++ runtime default 0. Defaults to None. temperature (float, optional): Controls the modulation of logits when sampling new tokens. It can have values >= 0.f. Defaults to None. @@ -379,6 +379,23 @@ def _validate(self): if self.temperature is not None and self.temperature < 0: raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") + # Top-p decay params follow the C++ warn-and-default policy (see + # topPSamplingLayer.cpp): out-of-range values are clamped to the runtime + # defaults rather than raising. Note top_p_min > top_p is intentionally + # allowed (C++ parity: the runtime top-p may rise toward top_p_min). + if self.top_p_decay is not None and (self.top_p_decay <= 0 or self.top_p_decay > 1): + logger.warning( + f"top_p_decay must be in (0, 1], got {self.top_p_decay}. " + f"Falling back to the default (1.0)." + ) + self.top_p_decay = None + if self.top_p_min is not None and (self.top_p_min <= 0 or self.top_p_min > 1): + logger.warning( + f"top_p_min must be in (0, 1], got {self.top_p_min}. " + f"Falling back to the default (1e-6)." + ) + self.top_p_min = None + if self.best_of is not None and self.best_of < self.n: raise ValueError(f"best_of ({self.best_of}) cannot be less than n ({self.n})") @@ -444,14 +461,42 @@ def params_imply_greedy_decoding( or temperature == 0 ) + # NB: Static, because downstream code only holds instances of + # bindings.SamplingConfig (not SamplingParams). Single source of truth for + # the "explicit greedy control" predicate shared with + # resolve_sampling_strategy (sampling_utils.py). + @staticmethod + def params_imply_explicit_greedy( + *, + temperature: Optional[float], + top_p: Optional[float], + top_k: Optional[int], + ) -> bool: + """Whether the request carries an explicit greedy control. + + Explicit means top_k == 1, top_p == 0.0, or temperature == 0, as opposed + to the implicit "all params unset" greedy default. + """ + return top_k == 1 or top_p == 0.0 or temperature == 0 + @property def _greedy_decoding(self) -> bool: - return self.params_imply_greedy_decoding( + if not self.params_imply_greedy_decoding( temperature=self.temperature, top_p=self.top_p, top_k=self.top_k, use_beam_search=self.use_beam_search, + ): + return False + # Keep this consistent with resolve_sampling_strategy: an active top_p_decay + # (set and < 1.0) routes to top-p sampling and is therefore NOT greedy, + # unless the request also carries an explicit greedy control, which is a + # deterministic intent that wins over decay. + decay_active = self.top_p_decay is not None and self.top_p_decay < 1.0 + explicitly_greedy = self.params_imply_explicit_greedy( + temperature=self.temperature, top_p=self.top_p, top_k=self.top_k ) + return explicitly_greedy or not decay_active @property def _need_return_context_logits(self) -> bool: diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index e2e0d1bc9402..82777e554479 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -58,7 +58,9 @@ TopK, TopKTopP, TopP, + TopPDecayOverride, UtilsSamplingParams, + resolve_sampling_strategy, ) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import SamplingConfig @@ -1793,6 +1795,7 @@ def _sample_grouped_strategies( generator: Optional[torch.Generator] = None, return_probs: bool, group_metadata: StrategyMetadata | None = None, + decay_override: Optional[TopPDecayOverride] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor] | float]: assert generator is sampler.get_generator(logits.device) if isinstance(group_key, tuple): @@ -1811,6 +1814,7 @@ def _sample_grouped_strategies( generator=generator, return_probs=return_probs, group_metadata=group_metadata, + decay_override=decay_override, ) ) return result @@ -2756,3 +2760,113 @@ def test_stale_conditional_and_unconditional_same_cell(self): logits = self._run_stale([req], [True], {0: 9}) assert self._banned_cols(logits[0]) == {2} assert not torch.isnan(logits).any() + + +class TestTopPDecay: + """Minimal functional guards for Top-P Decay in TorchSampler. + + Covers strategy routing, the post-sample runtime update (parity with the + C++ computeToppDecay recurrence; cases ported from + topPSamplingLayerTest.cpp), and per-request rejection of unsupported + combinations. + """ + + VOCAB_SIZE = 1000 + + @staticmethod + def _params(**kw) -> UtilsSamplingParams: + base = dict(temperature=None, top_p=None, top_k=None, use_beam_search=False) + base.update(kw) + return UtilsSamplingParams(**base) + + @staticmethod + def _make_sampler(*, max_draft_len=0): + return TorchSampler( + TorchSampler.Args( + max_seq_len=128, + max_draft_len=max_draft_len, + max_num_sequences=8, + max_beam_width=1, + max_total_draft_tokens=max_draft_len, + disable_overlap_scheduler=True, + ) + ) + + def test_strategy_routing(self): + # Active decay (set and < 1.0) forces a top-p-capable strategy even for + # an otherwise-greedy request (initial top-p defaults to 1.0), so the + # decayed runtime value can take effect on later steps. + s = resolve_sampling_strategy(self._params(top_p_decay=0.5), vocab_size=self.VOCAB_SIZE) + assert s[0] == "top_p" and s[1] == pytest.approx(1.0) + s = resolve_sampling_strategy( + self._params(top_k=50, top_p=0.9, top_p_decay=0.8), vocab_size=self.VOCAB_SIZE + ) + assert s[0] == "top_k_top_p" + # decay == 1.0 (the C++ default) is a no-op and does not activate... + s = resolve_sampling_strategy(self._params(top_p_decay=1.0), vocab_size=self.VOCAB_SIZE) + assert s is GREEDY + # ...and an explicit greedy control wins over an active decay. + s = resolve_sampling_strategy( + self._params(top_p_decay=0.5, top_k=1), vocab_size=self.VOCAB_SIZE + ) + assert s is GREEDY + + def test_runtime_update_parity(self): + # Post-sample update parity with the C++ computeToppDecay recurrence: + # runtime = initial if reset_id >= 0 and token == reset_id + # = max(runtime * decay, min) otherwise + sampler = self._make_sampler() + store = sampler.store.top_p_decay_store + configs = [ + dict(initial=0.8, decay=0.3, top_p_min=0.5, reset_id=2), # decay, then reset + dict(initial=0.2, decay=0.9, top_p_min=0.1, reset_id=-1), # plain decay, floored + dict(initial=0.3, decay=0.5, top_p_min=0.6, reset_id=-1), # min > initial: rises + ] + token_steps = [[1, 2, 3], [9, 9, 9], [9, 9, 9]] + slots = list(range(len(configs))) + for slot, cfg in zip(slots, configs): + sampler._top_p_decay_slots.add(slot) + store.runtime_top_p_decay_cuda[slot] = cfg["initial"] + store.initial_top_p_decay_cuda[slot] = cfg["initial"] + store.top_p_decay_cuda[slot] = cfg["decay"] + store.top_p_decay_min_cuda[slot] = cfg["top_p_min"] + store.top_p_decay_reset_ids_cuda[slot] = cfg["reset_id"] + store.is_top_p_decay_slot_cuda[slot] = True + + runtime = [cfg["initial"] for cfg in configs] + slots_cuda = torch.tensor(slots, dtype=torch.int64, device="cuda") + for step in range(3): + for slot in slots: + sampler.store.new_tokens[0, slot, 0] = token_steps[slot][step] + sampler._update_top_p_decay_after_sample( + new_tokens_cuda=sampler.store.new_tokens, sampled_slots_cuda=slots_cuda + ) + got = store.runtime_top_p_decay_cuda.cpu() + for slot, cfg in zip(slots, configs): + tok = token_steps[slot][step] + if cfg["reset_id"] >= 0 and tok == cfg["reset_id"]: + runtime[slot] = cfg["initial"] + else: + runtime[slot] = max(runtime[slot] * cfg["decay"], cfg["top_p_min"]) + assert got[slot].item() == pytest.approx(runtime[slot], abs=1e-6), (step, slot) + + def test_reject_speculative_draft_tokens(self): + # Decay + draft tokens through TorchSampler is rejected per-request at + # admission (validate_request), so only the offending request fails. + sampler = self._make_sampler(max_draft_len=4) + + def _request(params: SamplingParams, draft_tokens): + params._validate() + req = SimpleNamespace( + sampling_config=SamplingConfig(params._get_sampling_config()), + is_context_init_state=False, + py_sampling_strategy=None, + py_draft_tokens=draft_tokens, + ) + req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 + return cast(LlmRequest, req) + + with pytest.raises(ValueError, match="speculative"): + sampler.validate_request(_request(SamplingParams(top_p=0.9, top_p_decay=0.5), [1, 2])) + # Same request without decay is accepted. + sampler.validate_request(_request(SamplingParams(top_p=0.9), [1, 2])) From 2850f65b837c362e9676c945cabe02463ef6b218 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 16 Jul 2026 23:14:11 -0700 Subject: [PATCH 2/9] [TRTLLM-13231][chore] Drop the redundant reset-id gate in the decay update Sampled token ids are non-negative, so the negative reset_ids sentinel (-1, reset disabled) never matches without an explicit reset_id >= 0 test; compare directly, matching the legacy computeToppDecay kernel. Update the recurrence described in the op/kernel docs and the test reference accordingly. Addresses review feedback from @ixlmar. Signed-off-by: ZhaoyangWang --- .../kernels/samplerKernels/toppDecayKernels.cu | 6 ++++-- .../kernels/samplerKernels/toppDecayKernels.h | 8 ++++---- cpp/tensorrt_llm/thop/samplerOp.cpp | 2 +- tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py | 5 ++++- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 4 ++-- tests/unittest/_torch/sampler/test_torch_sampler.py | 7 ++++--- 6 files changed, 19 insertions(+), 13 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu index addcd038a33e..96ee83dbb823 100644 --- a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu +++ b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu @@ -40,9 +40,11 @@ __global__ void toppDecayUpdateKernel(float* __restrict__ runtimeTopP, float con return; } int32_t const tok = stepTokens[slot * stepTokenStride]; // gather in-kernel (strided new_tokens view) - int32_t const rid = resetIds[slot]; float updated; - if (rid >= 0 && tok == rid) + // Sampled token ids are non-negative, so a negative resetIds sentinel (-1, + // "reset disabled") never matches -- no explicit gate needed (parity with + // the legacy computeToppDecay kernel). + if (tok == resetIds[slot]) { updated = initialTopP[slot]; } diff --git a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h index 5e3eb2ead375..be6180abfe69 100644 --- a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h +++ b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h @@ -30,12 +30,12 @@ namespace kernels //! Unlike the legacy invokeComputeToppDecay (samplingTopPKernels.h), this variant //! gathers the sampled token in-kernel from a slot-indexed strided view of the //! new-tokens buffer (\p stepTokens with element stride \p stepTokenStride, i.e. -//! new_tokens[step, slot, beam] for a fixed step/beam), applies a gated reset -//! (reset_id < 0 disables the reset), and filters decay-active slots on-device -//! via \p isDecaySlot. +//! new_tokens[step, slot, beam] for a fixed step/beam) and filters decay-active +//! slots on-device via \p isDecaySlot. A negative resetIds entry (-1, "reset +//! disabled") never matches a sampled token, since token ids are non-negative. //! //! For each sampled row i with slot s = sampledSlots[i] where isDecaySlot[s]: -//! runtimeTopP[s] = (resetIds[s] >= 0 && stepTokens[s * stride] == resetIds[s]) +//! runtimeTopP[s] = (stepTokens[s * stride] == resetIds[s]) //! ? initialTopP[s] //! : max(runtimeTopP[s] * topPDecay[s], topPMin[s]) //! diff --git a/cpp/tensorrt_llm/thop/samplerOp.cpp b/cpp/tensorrt_llm/thop/samplerOp.cpp index ee39b5c44df8..6131509d476f 100644 --- a/cpp/tensorrt_llm/thop/samplerOp.cpp +++ b/cpp/tensorrt_llm/thop/samplerOp.cpp @@ -34,7 +34,7 @@ namespace torch_ext //////////////////////////////////////////////////////////////////////////////////////////////////////////// // Fused post-sample runtime top-p update for the TorchSampler Top-P Decay // feature. Applies, in place, for every sampled row whose slot is decay-active: -// runtime_top_p[slot] = (reset_id >= 0 && last_token == reset_id) +// runtime_top_p[slot] = (last_token == reset_id) // ? initial_top_p[slot] // : max(runtime_top_p[slot] * top_p_decay[slot], top_p_min[slot]) // The decay-active gate is applied on-device via is_decay_slot, so the hot path diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py index 70ece23a75d3..e5f3df9eb4ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py @@ -37,9 +37,12 @@ def top_p_decay_update( For each sampled row whose slot is decay-active (per ``is_decay_slot``):: - runtime_top_p[slot] = initial_top_p[slot] if reset_id >= 0 and token == reset_id + runtime_top_p[slot] = initial_top_p[slot] if token == reset_id = max(runtime_top_p[slot] * decay, top_p_min) otherwise + A negative ``reset_ids`` sentinel (-1, "reset disabled") never matches, + since sampled token ids are non-negative. + 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]``) -- diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 1b28a6bb8338..efe47ef325fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -4639,7 +4639,7 @@ def _update_top_p_decay_after_sample( """Apply the C++ ``computeToppDecay`` update for sampled decay-active slots. For each slot sampled this iteration that is decay-active: - runtime_p = initial_p if sampled_token == reset_id and reset_id >= 0 + runtime_p = initial_p if sampled_token == reset_id runtime_p = max(runtime_p * decay, top_p_min) otherwise Restricting to the sampled slots avoids reading stale new_tokens_cuda @@ -4656,7 +4656,7 @@ def _update_top_p_decay_after_sample( # sampled token in-kernel from the slot-indexed strided view of this # iteration's new-tokens buffer (decay is single-token-only, so the # token is at local step 0, beam 0), and applies - # runtime_p = (reset_id>=0 && tok==reset_id) ? initial_p + # runtime_p = (tok == reset_id) ? initial_p # : max(runtime_p*decay, top_p_min) # in place -- no separate gather/cast launches on the hot path. top_p_decay_update( diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 82777e554479..e7bf2834e190 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2812,8 +2812,9 @@ def test_strategy_routing(self): assert s is GREEDY def test_runtime_update_parity(self): - # Post-sample update parity with the C++ computeToppDecay recurrence: - # runtime = initial if reset_id >= 0 and token == reset_id + # Post-sample update parity with the C++ computeToppDecay recurrence + # (a negative reset_id never matches, since token ids are non-negative): + # runtime = initial if token == reset_id # = max(runtime * decay, min) otherwise sampler = self._make_sampler() store = sampler.store.top_p_decay_store @@ -2844,7 +2845,7 @@ def test_runtime_update_parity(self): got = store.runtime_top_p_decay_cuda.cpu() for slot, cfg in zip(slots, configs): tok = token_steps[slot][step] - if cfg["reset_id"] >= 0 and tok == cfg["reset_id"]: + if tok == cfg["reset_id"]: runtime[slot] = cfg["initial"] else: runtime[slot] = max(runtime[slot] * cfg["decay"], cfg["top_p_min"]) From 2e04c7ff207eec89d2f83c7c5e95c86226e13ae9 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 17 Jul 2026 00:15:53 -0700 Subject: [PATCH 3/9] [TRTLLM-13231][chore] Address review: encapsulation, naming, and comment cleanups - Extract all top-p-decay bookkeeping from setup_sampler_step into _setup_top_p_decay_for_new_requests; rename decay_params -> sampling_params. - Encapsulate TopPDecayStore allocation as TopPDecayStore.create(). - _retire_top_p_decay_slot: move the GENERATION_COMPLETE check to the call sites so it is not re-checked redundantly. - Narrow _build_top_p_runtime_override's strategy_key type to FlashInferGroupedStrategySampler.STRATEGY_KEY_TYPE. - Inline the float_tensor helper; generalize the SamplingConfig unwrap comment to cover all sampling fields; document that decay param range validation already lives in the executor::SamplingConfig constructor. - Tests: hoist the mock-request helper in TestTopPDecay. Addresses review feedback from @ixlmar. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/sampler/sampler.py | 182 ++++++++++-------- .../_torch/sampler/test_torch_sampler.py | 30 +-- 2 files changed, 118 insertions(+), 94 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index efe47ef325fb..a572fb03cb9a 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -478,14 +478,14 @@ def _get_max_beam_width(request: LlmRequest) -> int: def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: sampling_config = request.sampling_config - temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) - top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) - top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) - # These top-p-decay fields live on the C++ SamplingConfig as optional> + # These sampling fields live on the C++ SamplingConfig as optional> # (a shape designed for the batched TRT-LLM sampler); the torch sampler consumes - # them per request, so we unwrap the singleton list into a scalar here. When the + # them per request, so we unwrap the singleton lists into scalars here. When the # TRT-LLM sampler is removed, this SamplingConfig-based plumbing should be removed # too in favor of reading the values directly from the per-request params. + temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) + top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) + top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) top_p_decay = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_decay)) top_p_min = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_min)) top_p_reset_ids = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_p_reset_ids)) @@ -802,10 +802,6 @@ def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: return torch.empty(shape, dtype=torch.int, device=device) -def float_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: - return torch.empty(shape, dtype=torch.float32, device=device) - - @dataclass(kw_only=True, frozen=True) class _BatchedSamplingResult: # Original request indices for all requests (permuted due to batching by strategy): @@ -2290,6 +2286,21 @@ class TopPDecayStore: fused post-sample update kernel filter decay-active slots on the GPU, avoiding a host-side ``.tolist()`` / set intersection each step.""" + @classmethod + def create(cls, max_num_sequences: int) -> "TorchSampler.TopPDecayStore": + n = (max_num_sequences,) + return cls( + runtime_top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + initial_top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_min_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_reset_ids_cuda=torch.empty(n, dtype=torch.int, device="cuda"), + # The gate buffer IS the gate, so (unlike the others) it must start + # False: a slot that was never admitted as decay-active must read + # False. + is_top_p_decay_slot_cuda=torch.zeros(n, dtype=torch.bool, device="cuda"), + ) + @dataclass(kw_only=True) class Store: new_tokens: torch.Tensor @@ -2355,17 +2366,7 @@ def _create_store(self) -> Store: ) # Per-slot Top-P Decay runtime state (FlashInfer path). Allocated for all # sampler instances; only slots in self._top_p_decay_slots are ever read. - n = (self.max_num_sequences,) - top_p_decay_store = self.TopPDecayStore( - runtime_top_p_decay_cuda=float_tensor(n), - initial_top_p_decay_cuda=float_tensor(n), - top_p_decay_cuda=float_tensor(n), - top_p_decay_min_cuda=float_tensor(n), - top_p_decay_reset_ids_cuda=int_tensor(n), - # The gate buffer IS the gate, so (unlike the others) it must start - # False: a slot that was never admitted as decay-active must read False. - is_top_p_decay_slot_cuda=torch.zeros(n, dtype=torch.bool, device="cuda"), - ) + top_p_decay_store = self.TopPDecayStore.create(self.max_num_sequences) return self.Store( new_tokens=new_tokens, @@ -2533,6 +2534,12 @@ def _validate_top_p_decay_request(self, request: LlmRequest) -> None: request is failed individually instead of aborting the whole batch. """ params = _request_get_sampling_params(request) + # NB: value ranges need no re-check here. Every request enters through + # the executor::SamplingConfig constructor, which hard-validates + # top_p_decay in (0, 1], top_p_min in (0, 1] and top_p_reset_ids >= 0 + # (samplingConfig.cpp check* helpers) for all frontends. A reset id + # >= vocab_size is not checked anywhere but is semantically inert: it + # can never match a sampled token, i.e. it behaves as "reset disabled". if not top_p_decay_active(params): return if params.use_beam_search: @@ -3053,42 +3060,12 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: seq_slots: list[int] = [] # Used for beam search updates max_prompt_len: int = 0 - # Used for top-p-decay store updates (only populated for decay-active reqs) - decay_seq_slots: list[int] = [] - decay_initial_top_p: list[float] = [] - decay_top_p_decay: list[float] = [] - decay_top_p_min: list[float] = [] - decay_top_p_reset_ids: list[int] = [] - # Prepare finish reasons handler self._finish_reasons_handler.setup_new_request_handling() for request in new_requests: slot = request.py_seq_slot assert slot is not None seq_slots.append(slot) - # Drop any stale top-p-decay membership from a prior occupant of this - # slot; re-added below only if this request itself uses decay. - # (Unsupported decay combinations were already rejected per-request in - # validate_request at admission.) - self._top_p_decay_slots.discard(slot) - decay_params = _request_get_sampling_params(request) - if top_p_decay_active(decay_params): - self._top_p_decay_slots.add(slot) - decay_seq_slots.append(slot) - # Initial runtime top-p defaults to 1.0 when top_p is unset. - decay_initial_top_p.append( - decay_params.top_p if decay_params.top_p is not None else 1.0 - ) - # decay is guaranteed non-None and < 1.0 here (top_p_decay_active); - # min/reset fall back to the C++ runtime defaults when unset. - assert decay_params.top_p_decay is not None - decay_top_p_decay.append(decay_params.top_p_decay) - decay_top_p_min.append( - decay_params.top_p_min if decay_params.top_p_min is not None else 1e-6 - ) - decay_top_p_reset_ids.append( - decay_params.top_p_reset_ids if decay_params.top_p_reset_ids is not None else -1 - ) # update temp_data with this requests data self._finish_reasons_handler.prepare_for_new_request(request) @@ -3128,20 +3105,9 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: all_sampling_requests=new_requests + scheduled_requests.generation_requests, ) - # Clear the device decay gate for every newly-admitted slot (covers slot - # reuse: a slot previously decay-active but reused by a non-decay request - # must read False). Decay-active slots are then set True below. - decay_gate = self.store.top_p_decay_store.is_top_p_decay_slot_cuda - decay_gate.index_fill_(0, seq_slots_tensor_cuda_long, False) - - if decay_seq_slots: - self._update_top_p_decay_store_for_new_requests( - decay_seq_slots=decay_seq_slots, - initial_top_p=decay_initial_top_p, - top_p_decay=decay_top_p_decay, - top_p_min=decay_top_p_min, - top_p_reset_ids=decay_top_p_reset_ids, - ) + self._setup_top_p_decay_for_new_requests( + new_requests, new_seq_slots_cuda_long=seq_slots_tensor_cuda_long + ) if self._use_beam_search: beam_search_store = self.store.beam_search_store @@ -3153,6 +3119,65 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: max_prompt_len=max_prompt_len, ) + def _setup_top_p_decay_for_new_requests( + self, + new_requests: list[LlmRequest], + *, + new_seq_slots_cuda_long: torch.Tensor, + ) -> None: + """Refresh top-p-decay membership and per-slot buffers for admitted requests. + + Drops stale membership from prior occupants of the newly-admitted slots + (host set and device gate), then re-admits the decay-active requests and + initializes their per-slot store entries. Unsupported decay combinations + were already rejected per-request in validate_request at admission. + """ + # Clear the device decay gate for every newly-admitted slot (covers slot + # reuse: a slot previously decay-active but reused by a non-decay request + # must read False). Decay-active slots are then set True below. + decay_gate = self.store.top_p_decay_store.is_top_p_decay_slot_cuda + decay_gate.index_fill_(0, new_seq_slots_cuda_long, False) + + decay_seq_slots: list[int] = [] + initial_top_p: list[float] = [] + top_p_decay: list[float] = [] + top_p_min: list[float] = [] + top_p_reset_ids: list[int] = [] + for request in new_requests: + slot = request.py_seq_slot + assert slot is not None + self._top_p_decay_slots.discard(slot) + sampling_params = _request_get_sampling_params(request) + if not top_p_decay_active(sampling_params): + continue + self._top_p_decay_slots.add(slot) + decay_seq_slots.append(slot) + # Initial runtime top-p defaults to 1.0 when top_p is unset. + initial_top_p.append( + sampling_params.top_p if sampling_params.top_p is not None else 1.0 + ) + # decay is guaranteed non-None and < 1.0 here (top_p_decay_active); + # min/reset fall back to the C++ runtime defaults when unset. + assert sampling_params.top_p_decay is not None + top_p_decay.append(sampling_params.top_p_decay) + top_p_min.append( + sampling_params.top_p_min if sampling_params.top_p_min is not None else 1e-6 + ) + top_p_reset_ids.append( + sampling_params.top_p_reset_ids + if sampling_params.top_p_reset_ids is not None + else -1 + ) + + if decay_seq_slots: + self._update_top_p_decay_store_for_new_requests( + decay_seq_slots=decay_seq_slots, + initial_top_p=initial_top_p, + top_p_decay=top_p_decay, + top_p_min=top_p_min, + top_p_reset_ids=top_p_reset_ids, + ) + def _update_top_p_decay_store_for_new_requests( self, *, @@ -3870,7 +3895,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: for req_idx, req in enumerate(state.requests): if req.state == LlmRequestState.GENERATION_COMPLETE: - self._drop_top_p_decay_slot_if_finished(req) + self._retire_top_p_decay_slot(req) continue if req.py_beam_width > 1: @@ -3937,24 +3962,21 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: self._finish_reasons_handler.store.num_accepted_draft_tokens_host[ req.py_seq_slot ] = req.py_num_accepted_draft_tokens - self._drop_top_p_decay_slot_if_finished(req) + if req.state == LlmRequestState.GENERATION_COMPLETE: + self._retire_top_p_decay_slot(req) - def _drop_top_p_decay_slot_if_finished(self, req: LlmRequest) -> None: + def _retire_top_p_decay_slot(self, req: LlmRequest) -> None: """Retire a finished request's slot from the top-p-decay membership set. - Without this, `_top_p_decay_slots` would only shrink on slot reuse, so the - O(1) early-outs on the sampling hot path (post-sample update, override - build) would stay disabled forever once any decay request had run. The - device-side buffers need no cleanup: a freed slot is never sampled, and - slot reuse re-initializes them in setup_sampler_step. Requests that finish - outside the sampler (e.g. cancellation) are covered by the slot-reuse - cleanup instead. + Callers ensure ``req`` has finished. Without this, `_top_p_decay_slots` + would only shrink on slot reuse, so the O(1) early-outs on the sampling + hot path (post-sample update, override build) would stay disabled forever + once any decay request had run. The device-side buffers need no cleanup: + a freed slot is never sampled, and slot reuse re-initializes them in + setup_sampler_step. Requests that finish outside the sampler (e.g. + cancellation) are covered by the slot-reuse cleanup instead. """ - if ( - self._top_p_decay_slots - and req.state == LlmRequestState.GENERATION_COMPLETE - and req.py_seq_slot is not None - ): + if self._top_p_decay_slots and req.py_seq_slot is not None: self._top_p_decay_slots.discard(req.py_seq_slot) def _return_log_probs(self, requests: list[LlmRequest]) -> bool: @@ -4226,7 +4248,7 @@ def provision_bias_index() -> int: def _build_top_p_runtime_override( self, *, - strategy_key: Any, + strategy_key: FlashInferGroupedStrategySampler.STRATEGY_KEY_TYPE, group_req_indices: torch.Tensor, req_num_steps: torch.Tensor, seq_slots: torch.Tensor, diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index e7bf2834e190..90b0aec677a5 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2851,23 +2851,25 @@ def test_runtime_update_parity(self): runtime[slot] = max(runtime[slot] * cfg["decay"], cfg["top_p_min"]) assert got[slot].item() == pytest.approx(runtime[slot], abs=1e-6), (step, slot) + @staticmethod + def _mock_request(params: SamplingParams, *, draft_tokens=None): + params._validate() + req = SimpleNamespace( + sampling_config=SamplingConfig(params._get_sampling_config()), + is_context_init_state=False, + py_sampling_strategy=None, + py_draft_tokens=draft_tokens, + ) + req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 + return cast(LlmRequest, req) + def test_reject_speculative_draft_tokens(self): # Decay + draft tokens through TorchSampler is rejected per-request at # admission (validate_request), so only the offending request fails. sampler = self._make_sampler(max_draft_len=4) - - def _request(params: SamplingParams, draft_tokens): - params._validate() - req = SimpleNamespace( - sampling_config=SamplingConfig(params._get_sampling_config()), - is_context_init_state=False, - py_sampling_strategy=None, - py_draft_tokens=draft_tokens, - ) - req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 - return cast(LlmRequest, req) - with pytest.raises(ValueError, match="speculative"): - sampler.validate_request(_request(SamplingParams(top_p=0.9, top_p_decay=0.5), [1, 2])) + sampler.validate_request( + self._mock_request(SamplingParams(top_p=0.9, top_p_decay=0.5), draft_tokens=[1, 2]) + ) # Same request without decay is accepted. - sampler.validate_request(_request(SamplingParams(top_p=0.9), [1, 2])) + sampler.validate_request(self._mock_request(SamplingParams(top_p=0.9), draft_tokens=[1, 2])) From 173da42621bd1125c1bc7377cc18ecc0bb412637 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 17 Jul 2026 01:21:38 -0700 Subject: [PATCH 4/9] [TRTLLM-13231][chore] Route decay override via StrategyMetadata; reject invalid decay params Address review feedback from @ixlmar: - Pass the top-p-decay override through the existing StrategyMetadata mechanism (TopPDecayMetadata, mirroring BeamSearchMetadata) instead of a decay_override kwarg on sample_grouped_strategies; the override is applied inside the TopP*/TopKTopP* strategy impls via TopPDecayMixin (which declares _top_p, removing the getattr/setattr workaround). - Vectorize the per-STEP slot layout with torch.repeat_interleave; guard the single-token invariant in an if __debug__ block; use .item() over int(). - BREAKING: SamplingParams rejects out-of-range decay params (top_p_decay / top_p_min outside (0, 1], negative top_p_reset_ids) instead of warn-and-default, mirroring the executor::SamplingConfig hard checks. - Single-source the greedy/decay resolution: params_imply_top_p_decay_active and an extended params_imply_greedy_decoding(top_p_decay=...) on SamplingParams, reused by _greedy_decoding and resolve_sampling_strategy (decay predicate now evaluated lazily). - Consolidate the feature documentation: TopPDecayStore's docstring is the single source for semantics and the slot lifecycle; other sites keep their local contract and point back to it. Signed-off-by: ZhaoyangWang --- cpp/tensorrt_llm/thop/samplerOp.cpp | 17 +- docs/source/features/sampling.md | 3 +- .../_torch/pyexecutor/sampler/ops/trtllm.py | 11 +- .../_torch/pyexecutor/sampler/sampler.py | 209 +++++++++--------- .../pyexecutor/sampler/sampling_utils.py | 117 +++++----- tensorrt_llm/sampling_params.py | 105 +++++---- .../_torch/sampler/test_torch_sampler.py | 20 +- 7 files changed, 243 insertions(+), 239 deletions(-) diff --git a/cpp/tensorrt_llm/thop/samplerOp.cpp b/cpp/tensorrt_llm/thop/samplerOp.cpp index 6131509d476f..3ab902769e6f 100644 --- a/cpp/tensorrt_llm/thop/samplerOp.cpp +++ b/cpp/tensorrt_llm/thop/samplerOp.cpp @@ -32,13 +32,9 @@ namespace torch_ext { //////////////////////////////////////////////////////////////////////////////////////////////////////////// -// Fused post-sample runtime top-p update for the TorchSampler Top-P Decay -// feature. Applies, in place, for every sampled row whose slot is decay-active: -// runtime_top_p[slot] = (last_token == reset_id) -// ? initial_top_p[slot] -// : max(runtime_top_p[slot] * top_p_decay[slot], top_p_min[slot]) -// The decay-active gate is applied on-device via is_decay_slot, so the hot path -// needs no host-side .tolist() / set intersection. +// Fused post-sample runtime top-p update for the TorchSampler Top-P Decay feature. +// Semantics and update rule: see invokeToppDecayUpdate (kernels/samplerKernels/toppDecayKernels.h). +// This wrapper only validates the tensor contract and dispatches. void top_p_decay_update(th::Tensor runtime_top_p, th::Tensor initial_top_p, th::Tensor top_p_decay, th::Tensor top_p_min, th::Tensor reset_ids, th::Tensor is_decay_slot, th::Tensor step_tokens, th::Tensor sampled_slots) @@ -82,10 +78,9 @@ void top_p_decay_update(th::Tensor runtime_top_p, th::Tensor initial_top_p, th:: } //////////////////////////////////////////////////////////////////////////////////////////////////////////// -// Fused pre-sample per-row top-p gather for the Top-P Decay feature. Returns -// row_top_p[i] = is_decay_slot[slots[i]] ? runtime_top_p[slots[i]] : static_top_p[i] -// replacing the index_select(runtime) + index_select(gate) + where(static) chain -// with a single launch. +// Fused pre-sample per-row top-p gather for the Top-P Decay feature. +// Semantics: see invokeToppDecayGather (kernels/samplerKernels/toppDecayKernels.h). +// This wrapper only validates the tensor contract and dispatches. th::Tensor top_p_decay_gather( th::Tensor runtime_top_p, th::Tensor is_decay_slot, th::Tensor static_top_p, th::Tensor slots) { diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 437b68aaf1bf..72ad4a62ad41 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -105,7 +105,8 @@ llm.generate(["Hello, my name is", * 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). + 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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py index e5f3df9eb4ba..307308f0faf1 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py @@ -35,13 +35,10 @@ def top_p_decay_update( ) -> None: """Fused in-place update of ``runtime_top_p`` for the sampled decay slots. - For each sampled row whose slot is decay-active (per ``is_decay_slot``):: - - runtime_top_p[slot] = initial_top_p[slot] if token == reset_id - = max(runtime_top_p[slot] * decay, top_p_min) otherwise - - A negative ``reset_ids`` sentinel (-1, "reset disabled") never matches, - since sampled token ids are non-negative. + Applies the Top-P Decay recurrence (see ``TorchSampler.TopPDecayStore`` for + the feature-level semantics, or ``toppDecayKernels.h`` for the kernel + contract) 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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index a572fb03cb9a..fae43e11f4c9 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -98,7 +98,7 @@ GenericStrategyKeyType, Strategy, StrategyMetadata, - TopPDecayOverride, + TopPDecayMetadata, UtilsSamplingParams, get_rejected_indices, resolve_sampling_strategy, @@ -2263,12 +2263,45 @@ class LogProbsStore: @dataclass(kw_only=True) class TopPDecayStore: - """Per-slot runtime state for Top-P Decay (FlashInfer path). - - Each tensor has shape (max_num_sequences,). Only slots that are members - of ``TorchSampler._top_p_decay_slots`` hold meaningful values; all reads - and updates are gated by that set, so stale entries from a prior occupant - of a reused slot are never consumed. + """Per-slot runtime state for Top-P Decay -- the single source of truth + for the feature's semantics on the torch path. + + Semantics (matching the legacy C++ ``computeToppDecay`` kernel): after + every sampled token of a decay-active request:: + + runtime_top_p = initial_top_p if token == reset_id + = max(runtime_top_p * top_p_decay, top_p_min) otherwise + + A negative ``reset_ids`` sentinel (-1, "reset disabled") never matches, + since sampled token ids are non-negative. Decay is active iff + ``top_p_decay`` is set and < 1.0 (``SamplingParams. + params_imply_top_p_decay_active``); an active decay forces a + top-p-capable strategy even for an otherwise implicitly-greedy request + (initial top-p defaults to 1.0), while explicit greedy controls win. + Beam search and speculative draft tokens are rejected at admission + (``_validate_top_p_decay_request``); parameter ranges are enforced by + ``SamplingParams._validate`` / the executor::SamplingConfig constructor. + + Lifecycle per slot (each tensor has shape ``(max_num_sequences,)``): + + 1. Admission (``_setup_top_p_decay_for_new_requests``): membership is + cleared then re-set for the newly-admitted slots -- both the host-side + ``TorchSampler._top_p_decay_slots`` set (an O(1) hot-path early-out) + and its device mirror ``is_top_p_decay_slot_cuda`` (the gate the + fused kernels use, so the hot path needs no host-side filtering) -- + and the per-slot buffers are initialized. This clear-then-set also + covers slot reuse: stale entries from a prior occupant are never + consumed. + 2. Pre-sample (``_build_top_p_decay_metadata`` -> ``TopPDecayMetadata`` + -> ``TopPDecayMixin``): the per-row top-p fed to top_p / + top_k_top_p sampling is overridden with the decayed runtime value + for decay-active rows (fused gather, ``top_p_decay_gather``). + 3. Post-sample (``_update_top_p_decay_after_sample``): the recurrence + above is applied in place for the sampled decay-active slots (fused + update, ``top_p_decay_update``). + 4. Finish (``_retire_top_p_decay_slot``): the slot leaves the + membership set so the early-outs re-arm; the device buffers need no + cleanup (a freed slot is never sampled, reuse re-initializes it). """ runtime_top_p_decay_cuda: torch.Tensor @@ -2557,7 +2590,7 @@ def _validate_top_p_decay_request(self, request: LlmRequest) -> None: # usually not attached yet, so this check is best-effort. Two-model # speculation (the only source of such requests in TorchSampler) is # slated for removal, so in practice no speculative request reaches the - # decay path; a debug assert in _build_top_p_runtime_override guards the + # decay path; a debug assert in _build_top_p_decay_metadata guards the # invariant at sample time. if get_draft_token_length(request) > 0: raise ValueError( @@ -3125,7 +3158,8 @@ def _setup_top_p_decay_for_new_requests( *, new_seq_slots_cuda_long: torch.Tensor, ) -> None: - """Refresh top-p-decay membership and per-slot buffers for admitted requests. + """Refresh top-p-decay membership and per-slot buffers for admitted requests + (lifecycle step 1, see ``TopPDecayStore``). Drops stale membership from prior occupants of the newly-admitted slots (host set and device gate), then re-admits the decay-active requests and @@ -3703,6 +3737,7 @@ def _add_metadata_to_grouped_requests( *, seq_slots_cuda: torch.Tensor, seq_lens_cuda: torch.Tensor, + req_num_steps: torch.Tensor, ) -> dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata]: grouped_requests_with_metadata: dict[ RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata @@ -3712,6 +3747,7 @@ def _add_metadata_to_grouped_requests( num_requests = len(requests) for key, value in grouped_requests.items(): metadata_type = get_metadata_type_for_group_fn(key.strategy_key) + metadata: StrategyMetadata | None if metadata_type is BeamSearchMetadata: assert beam_search_store is not None assert seq_lens is not None, "seq_lens is required for beam search" @@ -3740,6 +3776,13 @@ def _add_metadata_to_grouped_requests( seq_offsets=beam_search_store.seq_offsets, beam_idx_arange=beam_search_store.beam_idx_arange, ) + elif metadata_type is TopPDecayMetadata: + metadata = self._build_top_p_decay_metadata( + group_req_indices=value.indices, + req_num_steps=req_num_steps, + seq_slots=seq_slots, + seq_slots_cuda=seq_slots_cuda, + ) elif metadata_type is None: metadata = None else: @@ -3966,15 +4009,13 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: self._retire_top_p_decay_slot(req) def _retire_top_p_decay_slot(self, req: LlmRequest) -> None: - """Retire a finished request's slot from the top-p-decay membership set. - - Callers ensure ``req`` has finished. Without this, `_top_p_decay_slots` - would only shrink on slot reuse, so the O(1) early-outs on the sampling - hot path (post-sample update, override build) would stay disabled forever - once any decay request had run. The device-side buffers need no cleanup: - a freed slot is never sampled, and slot reuse re-initializes them in - setup_sampler_step. Requests that finish outside the sampler (e.g. - cancellation) are covered by the slot-reuse cleanup instead. + """Retire a finished request's slot from the top-p-decay membership set + (lifecycle step 4, see ``TopPDecayStore``), so the O(1) hot-path + early-outs re-arm once decay traffic drains. + + Callers ensure ``req`` has finished. Requests that finish outside the + sampler (e.g. cancellation) are covered by the slot-reuse cleanup at + admission instead. """ if self._top_p_decay_slots and req.py_seq_slot is not None: self._top_p_decay_slots.discard(req.py_seq_slot) @@ -4245,26 +4286,24 @@ def provision_bias_index() -> int: # sharing). logits[logits_bias_mask_cuda] += biases_tensor_cuda - def _build_top_p_runtime_override( + def _build_top_p_decay_metadata( self, *, - strategy_key: FlashInferGroupedStrategySampler.STRATEGY_KEY_TYPE, group_req_indices: torch.Tensor, req_num_steps: torch.Tensor, seq_slots: torch.Tensor, seq_slots_cuda: torch.Tensor, - device: torch.device, - ) -> Optional[TopPDecayOverride]: - """Build the per-row runtime top-p override for a sampling group. - - Returns None unless this is a top-p / top_k_top_p group that may contain - a top-p-decay-active row. The override's ``slots`` tensor is aligned to - the group's per-STEP row order (matching group_strategies_per_step); - decay-active rows are single-token, while non-decay rows (possibly - multi-step draft rows) are gated out on-device by ``is_decay_slot``. + ) -> Optional[TopPDecayMetadata]: + """Build the Top-P Decay metadata for a top_p / top_k_top_p group. + + Lifecycle step 2, see ``TopPDecayStore``. Returns None when no request + currently uses decay. The metadata's ``slots`` tensor is aligned to the + group's per-STEP row order (matching group_strategies_per_step); + non-decay rows (possibly multi-step draft rows) are gated out on-device + by ``is_decay_slot``, so decay presence in the group is not checked + host-side: a group without decay rows samples every row with its static + top-p -- same result as returning None. """ - if strategy_key not in ("top_p", "top_k_top_p"): - return None if not self._top_p_decay_slots: return None store = self.store.top_p_decay_store @@ -4272,47 +4311,35 @@ def _build_top_p_runtime_override( # single-token, the per-STEP row order equals the per-request order, and # (group_req_indices being sorted ascending) a contiguous group's slots # are just a slice of seq_slots_cuda -- no host layout build and no H2D - # copy. Decay presence in the group is not checked host-side here: a - # group without decay rows is gated out on-device by is_decay_slot, so - # every row samples with its static top-p -- same result as returning - # None, minus the short-circuit. - first_req = int(group_req_indices[0]) - last_req = int(group_req_indices[-1]) + # copy. + first_req = group_req_indices[0].item() + last_req = group_req_indices[-1].item() + group_steps = req_num_steps[group_req_indices] if last_req - first_req + 1 == group_req_indices.size(0) and ( - int(req_num_steps[group_req_indices].max()) == 1 + group_steps.max().item() == 1 ): per_step_slots_cuda = seq_slots_cuda[first_req : last_req + 1] else: - group_seq_slots = seq_slots[group_req_indices].tolist() - group_steps = req_num_steps[group_req_indices].tolist() - # Build a per-STEP slot layout aligned with group_strategies_per_step - # (each request contributes req_num_steps rows). The per-row decay mask is - # not built here: it is gathered on-device from the is_decay_slot gate - # below. A decay-active row is always single-token: the only source of - # multi-step rows in TorchSampler is two-model speculation (slated for - # removal), and decay + draft tokens is rejected per-request at - # admission in validate_request -- so this is an internal invariant, - # not a user-reachable error path. - per_step_slots: list[int] = [] - has_decay = False - for slot, steps in zip(group_seq_slots, group_steps): - if slot in self._top_p_decay_slots: - has_decay = True - assert steps == 1, ( - "top_p_decay row with req_num_steps != 1; decay + draft tokens " - "should have been rejected at admission" - ) - per_step_slots.extend([slot] * steps) - if not has_decay: - return None - per_step_slots_cuda = torch.tensor( - per_step_slots, device="cpu", dtype=torch.int64, pin_memory=prefer_pinned() - ).to(device, non_blocking=True) - # The actual per-row gather (runtime top-p gated by is_decay_slot, the - # device mirror of _top_p_decay_slots kept in sync by setup_sampler_step's - # clear-then-set) happens in one fused launch inside - # _apply_top_p_runtime_override, where the static per-row top-p lives. - return TopPDecayOverride( + # Build the per-STEP slot layout (each request contributes + # req_num_steps rows). + group_seq_slots = seq_slots[group_req_indices] + if __debug__: + # Internal invariant (stripped under python -O): a decay-active + # row is always single-token -- the only source of multi-step + # rows in TorchSampler is two-model speculation (slated for + # removal), and decay + draft tokens is rejected per-request at + # admission in validate_request. + decay_row_steps = group_steps[ + torch.isin(group_seq_slots, torch.tensor(list(self._top_p_decay_slots))) + ] + assert decay_row_steps.numel() == 0 or decay_row_steps.max().item() == 1, ( + "top_p_decay row with req_num_steps != 1; decay + draft tokens " + "should have been rejected at admission" + ) + per_step_slots_cuda = torch.repeat_interleave(group_seq_slots.long(), group_steps).to( + seq_slots_cuda.device, non_blocking=True + ) + return TopPDecayMetadata( slots=per_step_slots_cuda, runtime_top_p=store.runtime_top_p_decay_cuda, is_decay_slot=store.is_top_p_decay_slot_cuda, @@ -4354,6 +4381,7 @@ def _sample_batched_by_strategy( get_metadata_type_for_group_fn=self._grouped_sampler_cls.get_metadata_type_for_group, seq_slots_cuda=seq_slots_cuda, seq_lens_cuda=seq_lens_cuda, + req_num_steps=req_num_steps, ) generator_cuda = self.get_generator(cuda_device) @@ -4476,24 +4504,6 @@ def _sample_batched_by_strategy( for _ in range(steps) ] - # For top-p / top_k_top_p groups, source the runtime (decayed) top-p - # for any decay-active rows from the per-slot store. Decay is single-token - # only, so per-step order matches per-request order here. Returns None - # on the non-FlashInfer path or when no decay row can be present, in - # which case the top-p-decay kwarg is omitted entirely (the Simple - # sampler backend does not accept it). - decay_override = self._build_top_p_runtime_override( - strategy_key=strategy_key, - group_req_indices=group_req_indices, - req_num_steps=req_num_steps, - seq_slots=seq_slots, - seq_slots_cuda=seq_slots_cuda, - device=group_logits_cuda.device, - ) - top_p_decay_kwargs = ( - {} if decay_override is None else {"decay_override": decay_override} - ) - group_next_tokens_cuda, group_softmax_cuda, group_temperature_cuda = ( self._grouped_sampler_cls.sample_grouped_strategies( strategy_key, @@ -4503,7 +4513,6 @@ def _sample_batched_by_strategy( return_probs=needs_probs, group_logit_indices=logit_indices_for_sampler, group_metadata=group_metadata, - **top_p_decay_kwargs, ) ) batch_next_tokens_offset_end = ( @@ -4658,29 +4667,19 @@ def _update_top_p_decay_after_sample( new_tokens_cuda: torch.Tensor, sampled_slots_cuda: torch.Tensor, ) -> None: - """Apply the C++ ``computeToppDecay`` update for sampled decay-active slots. - - For each slot sampled this iteration that is decay-active: - runtime_p = initial_p if sampled_token == reset_id - runtime_p = max(runtime_p * decay, top_p_min) otherwise + """Apply the post-sample decay recurrence for sampled decay-active slots. - Restricting to the sampled slots avoids reading stale new_tokens_cuda - for slots that were not scheduled this iteration. + See ``TopPDecayStore`` for the feature-level semantics (lifecycle + step 3). Restricting to the sampled slots avoids reading stale + new_tokens_cuda for slots that were not scheduled this iteration. """ - # Host-side O(1) early-out: skip entirely when no request uses decay - # (avoids launching the kernel). The per-slot decay-active filtering is - # otherwise done on-device via the is_decay_slot gate, so no .tolist() / - # set intersection is needed on the hot path. + # Host-side O(1) early-out: skip the kernel launch when no request uses + # decay; otherwise a single fused kernel gates on is_decay_slot + # on-device and gathers the sampled token in-kernel (decay is + # single-token-only, so the token is at local step 0, beam 0). if not self._top_p_decay_slots: return store = self.store.top_p_decay_store - # Single fused kernel: gates on is_decay_slot on-device, gathers the - # sampled token in-kernel from the slot-indexed strided view of this - # iteration's new-tokens buffer (decay is single-token-only, so the - # token is at local step 0, beam 0), and applies - # runtime_p = (tok == reset_id) ? initial_p - # : max(runtime_p*decay, top_p_min) - # in place -- no separate gather/cast launches on the hot path. top_p_decay_update( runtime_top_p=store.runtime_top_p_decay_cuda, initial_top_p=store.initial_top_p_decay_cuda, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index 6a7e6c2ef478..f24475954d4c 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -124,13 +124,16 @@ class UtilsSamplingParams: top_p_reset_ids: Optional[int] = None -@dataclass(frozen=True, kw_only=True) -class TopPDecayOverride: - """Per-group runtime top-p override for decay-active rows. +@dataclass(kw_only=True) +class TopPDecayMetadata(StrategyMetadata): + """Per-group runtime top-p override for Top-P Decay (attached to top_p / + top_k_top_p groups via the ``StrategyMetadata`` mechanism). ``slots`` maps each per-step group row to its sequence slot; the decayed per-row top-p is gathered on-device from the per-slot ``runtime_top_p`` store, gated by ``is_decay_slot`` (non-decay rows keep their static top-p). + Consumed by the TopP*/TopKTopP* strategy impls in ``sample()``. See + ``TorchSampler.TopPDecayStore`` for the feature-level semantics. """ slots: torch.Tensor @@ -144,11 +147,10 @@ class TopPDecayOverride: def top_p_decay_active(params: UtilsSamplingParams) -> bool: """Whether dynamic top-p decay is active for a request. - Decay is active iff ``top_p_decay`` is explicitly set and ``< 1.0``. A decay - of ``1.0`` (the C++ default) is a no-op, and ``top_p_min`` / ``top_p_reset_ids`` - alone do not activate dynamic behavior. + Delegates to the single-source predicate on SamplingParams; note that + ``top_p_min`` / ``top_p_reset_ids`` alone do not activate dynamic behavior. """ - return params.top_p_decay is not None and params.top_p_decay < 1.0 + return SamplingParams.params_imply_top_p_decay_active(params.top_p_decay) def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) -> Strategy: @@ -159,25 +161,17 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - top_p = params.top_p top_k = params.top_k - # When top-p decay is active, the runtime top-p can shrink below 1.0 on later - # steps even if the request starts at top_p == 1.0 / None. Force a top-p-capable - # strategy (defaulting the initial top-p to 1.0) so the decayed value takes effect. - # However, an EXPLICIT greedy control (see params_imply_explicit_greedy) is a - # deterministic intent that must be honored even with decay set -- only the - # implicit "all params unset" greedy case is overridden by decay. - decay_active = top_p_decay_active(params) - + # The greedy verdict (including the top-p-decay override of the implicit + # all-unset greedy default, and explicit greedy controls winning over decay) + # is single-sourced in SamplingParams.params_imply_greedy_decoding. if SamplingParams.params_imply_greedy_decoding( temperature=temperature, top_p=top_p, top_k=top_k, use_beam_search=use_beam_search, + top_p_decay=params.top_p_decay, ): - explicitly_greedy = SamplingParams.params_imply_explicit_greedy( - temperature=temperature, top_p=top_p, top_k=top_k - ) - if explicitly_greedy or not decay_active: - return GREEDY + return GREEDY # --- resolving default values # NB: not greedy, hence temperature != 0 if specified @@ -202,7 +196,7 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - # A decay-active request must go through a top-p-capable path even when its # initial top_p is 1.0, so the runtime top-p (sourced per-row at sample time) # can shrink the nucleus on later steps. - need_top_p = top_p < 1 or decay_active + need_top_p = top_p < 1 or top_p_decay_active(params) if need_top_p: if need_top_k: @@ -323,40 +317,6 @@ def _make_tensor(data: list[Any], dtype: torch.dtype, device: torch.device) -> t device=device, non_blocking=True ) - def _apply_top_p_runtime_override( - self, - decay_override: "TopPDecayOverride", - ) -> None: - """Override the per-row static top-p with the decayed top-p. - - Only decay-active rows (per the on-device ``is_decay_slot`` gate) are - overridden, so a group mixing top-p-decay and plain top-p requests - keeps each row's correct value. The overridden ``self._top_p`` tensor - then feeds both sampling and ``top_p_renorm_probs_op`` (so processed - logprobs match). Single fused launch (gather + gate + select). - - The caller (sample_grouped_strategies) is responsible for the - None-check; this method requires an actual override. - """ - # Only TopP*/TopKTopP* impls (which own a per-row ``_top_p`` tensor) - # receive a runtime override; access via getattr keeps this base-class - # helper independent of subclass attribute declarations. - static_top_p: torch.Tensor = getattr(self, "_top_p") - assert static_top_p.shape == decay_override.slots.shape, ( - static_top_p.shape, - decay_override.slots.shape, - ) - setattr( - self, - "_top_p", - top_p_decay_gather( - runtime_top_p=decay_override.runtime_top_p, - is_decay_slot=decay_override.is_decay_slot, - static_top_p=static_top_p, - slots=decay_override.slots, - ), - ) - @staticmethod def _prepare_logits_with_temperature( logits: torch.Tensor, @@ -430,6 +390,34 @@ def _sample_with_probs( new_tokens = cls._sample_from_probs(probs, generator=generator) return new_tokens, probs + class TopPDecayMixin: + """Mixed into the TopP*/TopKTopP* impls (the owners of a per-row + ``_top_p`` tensor) to consume ``TopPDecayMetadata``.""" + + _top_p: torch.Tensor + + def _maybe_apply_top_p_decay(self, group_metadata: Optional[StrategyMetadata]) -> None: + """Override the per-row static top-p with the decayed runtime top-p. + + Only decay-active rows (per the on-device ``is_decay_slot`` gate) are + overridden, so a group mixing top-p-decay and plain top-p requests + keeps each row's correct value. The overridden ``self._top_p`` tensor + then feeds both sampling and ``top_p_renorm_probs_op`` (so processed + logprobs match). Single fused launch (gather + gate + select). + """ + if not isinstance(group_metadata, TopPDecayMetadata): + return + assert self._top_p.shape == group_metadata.slots.shape, ( + self._top_p.shape, + group_metadata.slots.shape, + ) + self._top_p = top_p_decay_gather( + runtime_top_p=group_metadata.runtime_top_p, + is_decay_slot=group_metadata.is_decay_slot, + static_top_p=self._top_p, + slots=group_metadata.slots, + ) + class StrategyImplWithProbs(StrategyImpl): @override @classmethod @@ -458,7 +446,7 @@ def sample( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: return self._sample_greedy_with_probs(logits, group_logit_indices=group_logit_indices) - class TopKTopPWithProbs(StrategyImplWithProbs): + class TopKTopPWithProbs(TopPDecayMixin, StrategyImplWithProbs): def __init__(self, top_k: torch.Tensor, top_p: torch.Tensor, temperature: torch.Tensor): self._top_k = top_k self._top_p = top_p @@ -484,6 +472,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) return self._sample_with_probs( logits, group_logit_indices=group_logit_indices, @@ -526,7 +515,7 @@ def sample( generator=generator, ) - class TopPWithProbs(StrategyImplWithProbs): + class TopPWithProbs(TopPDecayMixin, StrategyImplWithProbs): def __init__(self, top_p: torch.Tensor, temperature: torch.Tensor): self._top_p = top_p self._temperature = temperature @@ -550,6 +539,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) return self._sample_with_probs( logits, group_logit_indices=group_logit_indices, @@ -618,7 +608,7 @@ def sample( logits = logits[group_logit_indices] return torch.argmax(logits, dim=-1), None - class TopKTopPSampleOnly(StrategyImplSampleOnly): + class TopKTopPSampleOnly(TopPDecayMixin, StrategyImplSampleOnly): def __init__(self, top_k: torch.Tensor, top_p: torch.Tensor, temperature: torch.Tensor): self._top_k = top_k self._top_p = top_p @@ -644,6 +634,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) logits = self._prepare_logits_with_temperature( logits, group_logit_indices, self._temperature ) @@ -689,7 +680,7 @@ def sample( check_nan=self._flashinfer_check_nans(probs), ), None - class TopPSampleOnly(StrategyImplSampleOnly): + class TopPSampleOnly(TopPDecayMixin, StrategyImplSampleOnly): def __init__(self, top_p: torch.Tensor, temperature: torch.Tensor): self._top_p = top_p self._temperature = temperature @@ -713,6 +704,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) probs = self._prepare_probs_with_temperature( logits, group_logit_indices, self._temperature ) @@ -839,6 +831,8 @@ def get_metadata_type_for_group( match strategy_key: case ("beam_search", _, _): return BeamSearchMetadata + case "top_p" | "top_k_top_p": + return TopPDecayMetadata case _: return None @@ -852,7 +846,6 @@ def sample_grouped_strategies( generator: Optional[torch.Generator] = None, return_probs: bool, group_metadata: StrategyMetadata | None = None, - decay_override: Optional[TopPDecayOverride] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: """Sample grouped strategies. @@ -902,10 +895,6 @@ def sample_grouped_strategies( else: assert group_logit_indices.size(0) == beam_width_in * len(strategies) strategy_impl = strategy_impl_cls.from_strategies(strategies, cuda_device=logits.device) - # For top-p / top_k_top_p groups, override the per-row static top-p with the - # decayed top-p on the decay-active rows (gated on-device). - if decay_override is not None: - strategy_impl._apply_top_p_runtime_override(decay_override) next_tokens, softmax = strategy_impl.sample( logits, group_logit_indices=group_logit_indices, diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index b19576964c16..e7fed5e2ced0 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -210,9 +210,9 @@ class SamplingParams: If neither temperature, top_p, nor top_k are specified, sampling is greedy. If temperature > 0 and/or top_k > 1 are specified, sampling will proceed accordingly and top_p will default to top_p = 1. Setting top_p = 0 should result in greedy sampling, but is currently disallowed in the backend. - top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. None means using C++ runtime default 1.e-6. Defaults to None. - top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. The token id which, when sampled, resets the decayed top-P to its initial value. None means using C++ runtime default -1 (which never matches a token). Defaults to None. - top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. None means using C++ runtime default 1.f. Defaults to None. + top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. Must be in (0, 1]; invalid values are rejected. None means using C++ runtime default 1.e-6. Defaults to None. + top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. The token id which, when sampled, resets the decayed top-P to its initial value. Must be >= 0; invalid values are rejected. None means using C++ runtime default -1 (which never matches a token). Defaults to None. + top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. Must be in (0, 1]; invalid values are rejected. None means using C++ runtime default 1.f. Defaults to None. seed (int, optional): Controls the random seed used by the random number generator in sampling. None means using C++ runtime default 0. Defaults to None. temperature (float, optional): Controls the modulation of logits when sampling new tokens. It can have values >= 0.f. Defaults to None. The value None is treated as "not specified" in the following. @@ -379,22 +379,19 @@ def _validate(self): if self.temperature is not None and self.temperature < 0: raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") - # Top-p decay params follow the C++ warn-and-default policy (see - # topPSamplingLayer.cpp): out-of-range values are clamped to the runtime - # defaults rather than raising. Note top_p_min > top_p is intentionally - # allowed (C++ parity: the runtime top-p may rise toward top_p_min). - if self.top_p_decay is not None and (self.top_p_decay <= 0 or self.top_p_decay > 1): - logger.warning( - f"top_p_decay must be in (0, 1], got {self.top_p_decay}. " - f"Falling back to the default (1.0)." - ) - self.top_p_decay = None - if self.top_p_min is not None and (self.top_p_min <= 0 or self.top_p_min > 1): - logger.warning( - f"top_p_min must be in (0, 1], got {self.top_p_min}. " - f"Falling back to the default (1e-6)." + # Top-p decay param ranges mirror the hard checks in the + # executor::SamplingConfig constructor (samplingConfig.cpp check* + # helpers); rejecting here gives a clear, early error instead of a + # RuntimeError from the C++ boundary. Note top_p_min > top_p is + # intentionally allowed (the runtime top-p may rise toward top_p_min). + if self.top_p_decay is not None and not 0.0 < self.top_p_decay <= 1.0: + raise ValueError(f"require 0 < top_p_decay <= 1, got top_p_decay={self.top_p_decay}") + if self.top_p_min is not None and not 0.0 < self.top_p_min <= 1.0: + raise ValueError(f"require 0 < top_p_min <= 1, got top_p_min={self.top_p_min}") + if self.top_p_reset_ids is not None and self.top_p_reset_ids < 0: + raise ValueError( + f"require top_p_reset_ids >= 0, got top_p_reset_ids={self.top_p_reset_ids}" ) - self.top_p_min = None if self.best_of is not None and self.best_of < self.n: raise ValueError(f"best_of ({self.best_of}) cannot be less than n ({self.n})") @@ -444,27 +441,23 @@ def _validate(self): if self.logprobs_simple_format and self.use_beam_search: raise ValueError("logprobs_simple_format is not supported with beam search") - # NB: Static, because downstream code only holds instances of - # bindings.SamplingConfig (not SamplingParams). + # NB: The predicates below are static because downstream code (e.g. + # sampling_utils.resolve_sampling_strategy) only holds instances of + # bindings.SamplingConfig (not SamplingParams). They are the single + # source of truth for the greedy / top-p-decay resolution shared by + # _greedy_decoding and the torch sampler. + @staticmethod - def params_imply_greedy_decoding( - *, - temperature: Optional[float], - top_p: Optional[float], - top_k: Optional[int], - use_beam_search: bool | None, - ): - return (not use_beam_search) and ( - (temperature is None and top_p is None and top_k is None) - or top_k == 1 - or top_p == 0.0 - or temperature == 0 - ) + def params_imply_top_p_decay_active(top_p_decay: Optional[float]) -> bool: + """Whether dynamic top-p decay is active. + + Active iff ``top_p_decay`` is explicitly set and ``< 1.0``; a decay of + ``1.0`` (the C++ default) is a no-op. Values outside ``(0, 1]`` are + rejected up front (_validate and the executor::SamplingConfig + constructor), so they never reach this predicate. + """ + return top_p_decay is not None and top_p_decay < 1.0 - # NB: Static, because downstream code only holds instances of - # bindings.SamplingConfig (not SamplingParams). Single source of truth for - # the "explicit greedy control" predicate shared with - # resolve_sampling_strategy (sampling_utils.py). @staticmethod def params_imply_explicit_greedy( *, @@ -479,24 +472,40 @@ def params_imply_explicit_greedy( """ return top_k == 1 or top_p == 0.0 or temperature == 0 + @staticmethod + def params_imply_greedy_decoding( + *, + temperature: Optional[float], + top_p: Optional[float], + top_k: Optional[int], + use_beam_search: bool | None, + top_p_decay: Optional[float] = None, + ) -> bool: + """Whether the parameters resolve to greedy decoding. + + An explicit greedy control always wins. The implicit "all params unset" + greedy default is overridden by an active top-p decay (which implies + top-p sampling so the decayed runtime top-p can take effect); callers + that do not support decay may omit ``top_p_decay``. + """ + if use_beam_search: + return False + if SamplingParams.params_imply_explicit_greedy( + temperature=temperature, top_p=top_p, top_k=top_k + ): + return True + implicitly_greedy = temperature is None and top_p is None and top_k is None + return implicitly_greedy and not SamplingParams.params_imply_top_p_decay_active(top_p_decay) + @property def _greedy_decoding(self) -> bool: - if not self.params_imply_greedy_decoding( + return self.params_imply_greedy_decoding( temperature=self.temperature, top_p=self.top_p, top_k=self.top_k, use_beam_search=self.use_beam_search, - ): - return False - # Keep this consistent with resolve_sampling_strategy: an active top_p_decay - # (set and < 1.0) routes to top-p sampling and is therefore NOT greedy, - # unless the request also carries an explicit greedy control, which is a - # deterministic intent that wins over decay. - decay_active = self.top_p_decay is not None and self.top_p_decay < 1.0 - explicitly_greedy = self.params_imply_explicit_greedy( - temperature=self.temperature, top_p=self.top_p, top_k=self.top_k + top_p_decay=self.top_p_decay, ) - return explicitly_greedy or not decay_active @property def _need_return_context_logits(self) -> bool: diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 90b0aec677a5..be3a94efabe8 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -58,7 +58,6 @@ TopK, TopKTopP, TopP, - TopPDecayOverride, UtilsSamplingParams, resolve_sampling_strategy, ) @@ -1795,7 +1794,6 @@ def _sample_grouped_strategies( generator: Optional[torch.Generator] = None, return_probs: bool, group_metadata: StrategyMetadata | None = None, - decay_override: Optional[TopPDecayOverride] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor] | float]: assert generator is sampler.get_generator(logits.device) if isinstance(group_key, tuple): @@ -1814,7 +1812,6 @@ def _sample_grouped_strategies( generator=generator, return_probs=return_probs, group_metadata=group_metadata, - decay_override=decay_override, ) ) return result @@ -2863,6 +2860,23 @@ def _mock_request(params: SamplingParams, *, draft_tokens=None): req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 return cast(LlmRequest, req) + @pytest.mark.parametrize( + "bad_kwargs", + [ + {"top_p_decay": 1.5}, + {"top_p_decay": -0.5}, + {"top_p_decay": 0.0}, + {"top_p_min": 0.0}, + {"top_p_min": 1.5}, + {"top_p_reset_ids": -1}, + ], + ) + def test_out_of_range_decay_params_rejected(self, bad_kwargs): + # Out-of-range decay params raise (mirroring the executor::SamplingConfig + # constructor's hard checks) instead of the former warn-and-default. + with pytest.raises(ValueError): + SamplingParams(**bad_kwargs) + def test_reject_speculative_draft_tokens(self): # Decay + draft tokens through TorchSampler is rejected per-request at # admission (validate_request), so only the offending request fails. From 8e2dea1bf6e92af34f25b603a3001328bc3fed52 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 17 Jul 2026 01:46:56 -0700 Subject: [PATCH 5/9] [TRTLLM-13231][fix] Pass seq_slots_cuda in the unbatch test caller _unbatch_sampling_results gained a required seq_slots_cuda kwarg (the resident device copy used by the post-sample top-p-decay update), but the direct caller in test_unbatch_sampling_results was not updated. Precompute the device copy outside the no-sync region. Full test_torch_sampler.py now passes (205/205) in a matching container. Signed-off-by: ZhaoyangWang --- tests/unittest/_torch/sampler/test_torch_sampler.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index be3a94efabe8..892b93357fb0 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2578,6 +2578,11 @@ class UutResultWrapper: result: Optional[UutResult] = None res = UutResultWrapper() + # Precomputed outside the no-sync region (mirrors the production + # resident device copy of seq_slots). + seq_slots_tensor_cuda = ( + seq_slots_tensor.to(torch.int64).pin_memory().to("cuda", non_blocking=True) + ) def _uut(res=res): new_tokens_host = sampler._unbatch_sampling_results( @@ -2585,6 +2590,7 @@ def _uut(res=res): new_tokens_cuda=new_tokens_cuda, req_num_generated_tokens=req_num_steps, seq_slots=seq_slots_tensor, + seq_slots_cuda=seq_slots_tensor_cuda, ) res.result = UutResult(new_tokens_host=new_tokens_host) From 17451b1a4f1042d7d4ce86d31043a0fac675795c Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Jul 2026 00:23:55 -0700 Subject: [PATCH 6/9] [TRTLLM-13231][refactor] Replace the top-p-decay CUDA kernels with torch.compile'd ops Per review discussion (@ixlmar): keep the new sampler kernel-free. The two decay ops (pre-sample per-row gather, post-sample in-place update) are now native torch functions fused with torch.compile in sampler/ops/top_p_decay.py; the custom CUDA kernels, the thop op wrappers, and their build wiring are removed. Compile configuration and rationale (benchmarked on H200, bs=32, per-step sampler cost beyond a fake model forward): - mode=max-autotune-no-cudagraphs: cudagraphs is unsafe here (in-place per-slot state, output consumed outside the compiled region) and the cudagraphs-enabled modes were either slower or produced overwritten outputs. - torch._dynamo.mark_dynamic on the batch-varying dims avoids recompilation as batch composition changes. - With a >=2ms forward, this matches the removed fused kernels (~141us vs ~142us per step, both ~35us ahead of eager); with a ~1ms forward it costs ~40us/step over eager (~90us over the fused kernels), the fixed compiled-entry cost being exposed. The kernel-free maintainability trade-off was judged worth it for the non-critical-path decay feature. - Compilation is lazy: first decay-active request pays ~1s; non-decay workloads never trigger it. Signed-off-by: ZhaoyangWang --- .../samplerKernels/toppDecayKernels.cu | 101 -------------- .../kernels/samplerKernels/toppDecayKernels.h | 59 -------- cpp/tensorrt_llm/thop/CMakeLists.txt | 1 - cpp/tensorrt_llm/thop/samplerOp.cpp | 132 ------------------ .../pyexecutor/sampler/ops/top_p_decay.py | 127 +++++++++++++++++ .../_torch/pyexecutor/sampler/ops/trtllm.py | 78 ----------- .../_torch/pyexecutor/sampler/sampler.py | 8 +- .../pyexecutor/sampler/sampling_utils.py | 4 +- 8 files changed, 133 insertions(+), 377 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu delete mode 100644 cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h delete mode 100644 cpp/tensorrt_llm/thop/samplerOp.cpp create mode 100644 tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py delete mode 100644 tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py diff --git a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu deleted file mode 100644 index 96ee83dbb823..000000000000 --- a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cu +++ /dev/null @@ -1,101 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -namespace -{ - -__global__ void toppDecayUpdateKernel(float* __restrict__ runtimeTopP, float const* __restrict__ initialTopP, - float const* __restrict__ topPDecay, float const* __restrict__ topPMin, int32_t const* __restrict__ resetIds, - bool const* __restrict__ isDecaySlot, int32_t const* __restrict__ stepTokens, int64_t stepTokenStride, - int64_t const* __restrict__ sampledSlots, int32_t numSampled) -{ - int32_t const i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= numSampled) - { - return; - } - int64_t const slot = sampledSlots[i]; - if (!isDecaySlot[slot]) // on-device decay gate - { - return; - } - int32_t const tok = stepTokens[slot * stepTokenStride]; // gather in-kernel (strided new_tokens view) - float updated; - // Sampled token ids are non-negative, so a negative resetIds sentinel (-1, - // "reset disabled") never matches -- no explicit gate needed (parity with - // the legacy computeToppDecay kernel). - if (tok == resetIds[slot]) - { - updated = initialTopP[slot]; - } - else - { - updated = fmaxf(runtimeTopP[slot] * topPDecay[slot], topPMin[slot]); - } - runtimeTopP[slot] = updated; -} - -__global__ void toppDecayGatherKernel(float* __restrict__ rowTopP, float const* __restrict__ runtimeTopP, - bool const* __restrict__ isDecaySlot, float const* __restrict__ staticTopP, int64_t const* __restrict__ slots, - int32_t numRows) -{ - int32_t const i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= numRows) - { - return; - } - int64_t const slot = slots[i]; - rowTopP[i] = isDecaySlot[slot] ? runtimeTopP[slot] : staticTopP[i]; -} - -} // namespace - -void invokeToppDecayUpdate(float* runtimeTopP, float const* initialTopP, float const* topPDecay, float const* topPMin, - int32_t const* resetIds, bool const* isDecaySlot, int32_t const* stepTokens, int64_t stepTokenStride, - int64_t const* sampledSlots, int32_t numSampled, cudaStream_t stream) -{ - if (numSampled == 0) - { - return; - } - constexpr int32_t kBlock = 256; - int32_t const grid = (numSampled + kBlock - 1) / kBlock; - toppDecayUpdateKernel<<>>(runtimeTopP, initialTopP, topPDecay, topPMin, resetIds, - isDecaySlot, stepTokens, stepTokenStride, sampledSlots, numSampled); -} - -void invokeToppDecayGather(float* rowTopP, float const* runtimeTopP, bool const* isDecaySlot, float const* staticTopP, - int64_t const* slots, int32_t numRows, cudaStream_t stream) -{ - if (numRows == 0) - { - return; - } - constexpr int32_t kBlock = 256; - int32_t const grid = (numRows + kBlock - 1) / kBlock; - toppDecayGatherKernel<<>>(rowTopP, runtimeTopP, isDecaySlot, staticTopP, slots, numRows); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h b/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h deleted file mode 100644 index be6180abfe69..000000000000 --- a/cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "tensorrt_llm/common/config.h" -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -//! \brief Fused post-sample runtime top-p update for the PyTorch TorchSampler's -//! Top-P Decay feature. One thread per sampled row; the decay-active gate is -//! applied on-device via \p isDecaySlot so no host-side filtering is needed. -//! -//! Unlike the legacy invokeComputeToppDecay (samplingTopPKernels.h), this variant -//! gathers the sampled token in-kernel from a slot-indexed strided view of the -//! new-tokens buffer (\p stepTokens with element stride \p stepTokenStride, i.e. -//! new_tokens[step, slot, beam] for a fixed step/beam) and filters decay-active -//! slots on-device via \p isDecaySlot. A negative resetIds entry (-1, "reset -//! disabled") never matches a sampled token, since token ids are non-negative. -//! -//! For each sampled row i with slot s = sampledSlots[i] where isDecaySlot[s]: -//! runtimeTopP[s] = (stepTokens[s * stride] == resetIds[s]) -//! ? initialTopP[s] -//! : max(runtimeTopP[s] * topPDecay[s], topPMin[s]) -//! -//! All per-slot arrays are length numSlots; \p sampledSlots is length numSampled. -//! \p runtimeTopP is updated in place. -void invokeToppDecayUpdate(float* runtimeTopP, float const* initialTopP, float const* topPDecay, float const* topPMin, - int32_t const* resetIds, bool const* isDecaySlot, int32_t const* stepTokens, int64_t stepTokenStride, - int64_t const* sampledSlots, int32_t numSampled, cudaStream_t stream); - -//! \brief Fused pre-sample per-row top-p gather for the Top-P Decay feature. -//! Replaces the eager index_select(runtime) + index_select(gate) + where(static) -//! chain with a single launch: -//! rowTopP[i] = isDecaySlot[slots[i]] ? runtimeTopP[slots[i]] : staticTopP[i] -//! \p rowTopP / \p staticTopP / \p slots are length numRows (per-step group rows); -//! \p runtimeTopP / \p isDecaySlot are per-slot arrays. -void invokeToppDecayGather(float* rowTopP, float const* runtimeTopP, bool const* isDecaySlot, float const* staticTopP, - int64_t const* slots, int32_t numRows, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 61912effeba2..b95426f580fe 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -113,7 +113,6 @@ add_library( IndexerKCacheGatherOp.cpp IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp - samplerOp.cpp mlaRopeInplaceOp.cpp ncclCommunicatorOp.cpp allocateOutput.cpp diff --git a/cpp/tensorrt_llm/thop/samplerOp.cpp b/cpp/tensorrt_llm/thop/samplerOp.cpp deleted file mode 100644 index 3ab902769e6f..000000000000 --- a/cpp/tensorrt_llm/thop/samplerOp.cpp +++ /dev/null @@ -1,132 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -// Torch custom ops for the PyTorch backend sampler (TorchSampler). Aggregates -// sampler-related ops here so follow-up work can add new sampler ops alongside -// existing ones, mirroring how specDecOp.cpp groups the speculative-decoding ops. - -#include "tensorrt_llm/common/opUtils.h" -#include "tensorrt_llm/kernels/samplerKernels/toppDecayKernels.h" -#include "tensorrt_llm/runtime/torchUtils.h" - -namespace th = torch; -namespace tk = tensorrt_llm::kernels; - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////////////// -// Fused post-sample runtime top-p update for the TorchSampler Top-P Decay feature. -// Semantics and update rule: see invokeToppDecayUpdate (kernels/samplerKernels/toppDecayKernels.h). -// This wrapper only validates the tensor contract and dispatches. -void top_p_decay_update(th::Tensor runtime_top_p, th::Tensor initial_top_p, th::Tensor top_p_decay, - th::Tensor top_p_min, th::Tensor reset_ids, th::Tensor is_decay_slot, th::Tensor step_tokens, - th::Tensor sampled_slots) -{ - TORCH_CHECK(runtime_top_p.is_cuda() && initial_top_p.is_cuda() && top_p_decay.is_cuda() && top_p_min.is_cuda() - && reset_ids.is_cuda() && is_decay_slot.is_cuda() && step_tokens.is_cuda() && sampled_slots.is_cuda(), - "all top_p_decay_update tensors must be CUDA tensors"); - - TORCH_CHECK(runtime_top_p.scalar_type() == th::kFloat32, "runtime_top_p must be float32"); - TORCH_CHECK(initial_top_p.scalar_type() == th::kFloat32, "initial_top_p must be float32"); - TORCH_CHECK(top_p_decay.scalar_type() == th::kFloat32, "top_p_decay must be float32"); - TORCH_CHECK(top_p_min.scalar_type() == th::kFloat32, "top_p_min must be float32"); - TORCH_CHECK(reset_ids.scalar_type() == th::kInt32, "reset_ids must be int32"); - TORCH_CHECK(is_decay_slot.scalar_type() == th::kBool, "is_decay_slot must be bool"); - TORCH_CHECK(step_tokens.scalar_type() == th::kInt32, "step_tokens must be int32"); - TORCH_CHECK(sampled_slots.scalar_type() == th::kInt64, "sampled_slots must be int64"); - - // step_tokens is a slot-indexed 1-D view of the new-tokens buffer for a fixed - // step/beam (new_tokens[step, :, beam]); it is strided, not contiguous. The - // kernel gathers tokens through its element stride, avoiding a separate - // gather + cast launch on the hot path. - TORCH_CHECK(runtime_top_p.is_contiguous() && initial_top_p.is_contiguous() && top_p_decay.is_contiguous() - && top_p_min.is_contiguous() && reset_ids.is_contiguous() && is_decay_slot.is_contiguous() - && sampled_slots.is_contiguous(), - "all top_p_decay_update tensors (except step_tokens) must be contiguous"); - TORCH_CHECK(step_tokens.dim() == 1, "step_tokens must be a 1-D (strided) view"); - - auto const numSlots = runtime_top_p.size(0); - TORCH_CHECK(initial_top_p.size(0) == numSlots && top_p_decay.size(0) == numSlots && top_p_min.size(0) == numSlots - && reset_ids.size(0) == numSlots && is_decay_slot.size(0) == numSlots, - "all per-slot tensors must have the same length (max_num_sequences)"); - TORCH_CHECK(step_tokens.size(0) >= numSlots, "step_tokens must cover all slots"); - - auto const numSampled = sampled_slots.size(0); - - auto stream = at::cuda::getCurrentCUDAStream(runtime_top_p.get_device()); - tk::invokeToppDecayUpdate(runtime_top_p.data_ptr(), initial_top_p.data_ptr(), - top_p_decay.data_ptr(), top_p_min.data_ptr(), reset_ids.data_ptr(), - is_decay_slot.data_ptr(), step_tokens.data_ptr(), step_tokens.stride(0), - sampled_slots.data_ptr(), static_cast(numSampled), stream); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////// -// Fused pre-sample per-row top-p gather for the Top-P Decay feature. -// Semantics: see invokeToppDecayGather (kernels/samplerKernels/toppDecayKernels.h). -// This wrapper only validates the tensor contract and dispatches. -th::Tensor top_p_decay_gather( - th::Tensor runtime_top_p, th::Tensor is_decay_slot, th::Tensor static_top_p, th::Tensor slots) -{ - TORCH_CHECK(runtime_top_p.is_cuda() && is_decay_slot.is_cuda() && static_top_p.is_cuda() && slots.is_cuda(), - "all top_p_decay_gather tensors must be CUDA tensors"); - TORCH_CHECK(runtime_top_p.scalar_type() == th::kFloat32, "runtime_top_p must be float32"); - TORCH_CHECK(is_decay_slot.scalar_type() == th::kBool, "is_decay_slot must be bool"); - TORCH_CHECK(static_top_p.scalar_type() == th::kFloat32, "static_top_p must be float32"); - TORCH_CHECK(slots.scalar_type() == th::kInt64, "slots must be int64"); - TORCH_CHECK(runtime_top_p.is_contiguous() && is_decay_slot.is_contiguous() && static_top_p.is_contiguous() - && slots.is_contiguous(), - "all top_p_decay_gather tensors must be contiguous"); - TORCH_CHECK(runtime_top_p.size(0) == is_decay_slot.size(0), "per-slot tensors must have the same length"); - auto const numRows = slots.size(0); - TORCH_CHECK(static_top_p.size(0) == numRows, "static_top_p and slots must have the same length"); - // NB: every slots[i] must lie in [0, runtime_top_p.size(0)); the kernel indexes the per-slot - // arrays with it unchecked (values live on device, so validating here would force a sync). - // The caller (TorchSampler) guarantees this: slots come from seq_slots, which is bounded by - // max_num_sequences -- the allocation size of the per-slot store tensors. - - auto row_top_p = th::empty_like(static_top_p); - auto stream = at::cuda::getCurrentCUDAStream(runtime_top_p.get_device()); - tk::invokeToppDecayGather(row_top_p.data_ptr(), runtime_top_p.data_ptr(), - is_decay_slot.data_ptr(), static_top_p.data_ptr(), slots.data_ptr(), - static_cast(numRows), stream); - return row_top_p; -} - -} // end namespace torch_ext - -TRTLLM_NAMESPACE_END - -//////////////////////////////////////////////////////////////////////////////////////////////////////////// - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "top_p_decay_update(Tensor(a!) runtime_top_p, Tensor initial_top_p, Tensor top_p_decay, Tensor top_p_min, " - "Tensor reset_ids, Tensor is_decay_slot, Tensor step_tokens, Tensor sampled_slots) -> ()"); - m.def( - "top_p_decay_gather(Tensor runtime_top_p, Tensor is_decay_slot, Tensor static_top_p, Tensor slots) " - "-> Tensor"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("top_p_decay_update", &tensorrt_llm::torch_ext::top_p_decay_update); - m.impl("top_p_decay_gather", &tensorrt_llm::torch_ext::top_p_decay_gather); -} diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py new file mode 100644 index 000000000000..1aac8d8ab218 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Top-P Decay ops for TorchSampler, fused via ``torch.compile``. + +These two ops sit on a host-launch-bound path (per-step work of a few dozen +elements per row), so the eager op sequences are fused with Inductor to keep +the per-step launch count low; the generated Triton kernels match the GPU-side +cost of hand-written fused kernels once the model forward is long enough to +hide the compiled-function entry. + +Compile configuration (benchmarked against eager and other modes on H200): + +- ``mode="max-autotune-no-cudagraphs"``: CUDA graphs are deliberately + disabled -- the update mutates persistent per-slot state in place and the + gather's output is consumed outside the compiled region, both of which are + unsafe with cudagraph static buffers (outputs get overwritten by subsequent + replays). +- ``torch._dynamo.mark_dynamic`` on the batch-varying dimensions, so per-step + changes in batch composition do not trigger recompilation while the + remaining dimensions stay specialized. + +Compilation happens lazily on the first decay-active request (roughly a +second); non-decay workloads never trigger it. +""" + +import torch + + +@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) + + +@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) + + +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) + _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, + ) + + +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 _top_p_decay_gather_impl(runtime_top_p, is_decay_slot, static_top_p, slots) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py deleted file mode 100644 index 307308f0faf1..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Thin Python wrappers over the trtllm C++ sampling ops used by TorchSampler. - -These forward to ``torch.ops.trtllm.*`` custom ops (registered from -``cpp/tensorrt_llm/thop``); the wrappers exist to give the sampler a stable, -keyword-only Python surface decoupled from the raw op-registration names. -""" - -import torch - - -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, or ``toppDecayKernels.h`` for the kernel - contract) 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]``) -- - the kernel gathers each sampled token in-kernel; ``sampled_slots`` is 1-D of - length ``num_sampled`` (this iteration's rows). ``runtime_top_p`` is mutated - in place; nothing is returned. - """ - torch.ops.trtllm.top_p_decay_update( - runtime_top_p, - initial_top_p, - top_p_decay, - top_p_min, - reset_ids, - is_decay_slot, - step_tokens, - sampled_slots, - ) - - -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 - - replacing the eager ``index_select`` x2 + ``where`` chain with one launch. - """ - return torch.ops.trtllm.top_p_decay_gather(runtime_top_p, is_decay_slot, static_top_p, slots) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index fae43e11f4c9..3b1e5eda0078 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -88,7 +88,7 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests -from .ops.trtllm import top_p_decay_update +from .ops.top_p_decay import top_p_decay_update from .sampling_utils import ( BEAM_SEARCH_PAD_TOKEN, GREEDY, @@ -2316,7 +2316,7 @@ class TopPDecayStore: """Per-slot reset token id (< 0 never matches a sampled token).""" is_top_p_decay_slot_cuda: torch.Tensor """Per-slot bool gate (device mirror of ``_top_p_decay_slots``). Lets the - fused post-sample update kernel filter decay-active slots on the GPU, + fused post-sample update op filter decay-active slots on the GPU, avoiding a host-side ``.tolist()`` / set intersection each step.""" @classmethod @@ -4674,8 +4674,8 @@ def _update_top_p_decay_after_sample( new_tokens_cuda for slots that were not scheduled this iteration. """ # Host-side O(1) early-out: skip the kernel launch when no request uses - # decay; otherwise a single fused kernel gates on is_decay_slot - # on-device and gathers the sampled token in-kernel (decay is + # decay; otherwise a single fused (torch.compile) op gates on + # is_decay_slot on-device and gathers the sampled token in place (decay is # single-token-only, so the token is at local step 0, beam 0). if not self._top_p_decay_slots: return diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index f24475954d4c..b3a9f1bfe6c2 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -39,7 +39,7 @@ top_p_renorm_probs_op, top_p_sampling_from_probs_op, ) -from tensorrt_llm._torch.pyexecutor.sampler.ops.trtllm import top_p_decay_gather +from tensorrt_llm._torch.pyexecutor.sampler.ops.top_p_decay import top_p_decay_gather from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import ( GREEDY_TEMPERATURE_THRESHOLD, BeamSearchMetadata, @@ -403,7 +403,7 @@ def _maybe_apply_top_p_decay(self, group_metadata: Optional[StrategyMetadata]) - overridden, so a group mixing top-p-decay and plain top-p requests keeps each row's correct value. The overridden ``self._top_p`` tensor then feeds both sampling and ``top_p_renorm_probs_op`` (so processed - logprobs match). Single fused launch (gather + gate + select). + logprobs match). Fused via torch.compile (gather + gate + select). """ if not isinstance(group_metadata, TopPDecayMetadata): return From 645df100dfdcbf9dd10ffc49b0a075ded687c079 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Jul 2026 00:28:40 -0700 Subject: [PATCH 7/9] [TRTLLM-13231][chore] Move the top-p-decay fused ops into ops.vanilla.Fusions They are pure-torch compiled ops; Fusions in ops/vanilla.py is the existing home for torch.compile'd fused helpers (same impl + mark_dynamic wrapper pattern as gather_log_softmax), so the standalone ops/top_p_decay.py module is folded into it and call sites use Fusions.top_p_decay_{gather,update}. Signed-off-by: ZhaoyangWang --- .../pyexecutor/sampler/ops/top_p_decay.py | 127 ------------------ .../_torch/pyexecutor/sampler/ops/vanilla.py | 102 ++++++++++++++ .../_torch/pyexecutor/sampler/sampler.py | 3 +- .../pyexecutor/sampler/sampling_utils.py | 3 +- 4 files changed, 104 insertions(+), 131 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py deleted file mode 100644 index 1aac8d8ab218..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Top-P Decay ops for TorchSampler, fused via ``torch.compile``. - -These two ops sit on a host-launch-bound path (per-step work of a few dozen -elements per row), so the eager op sequences are fused with Inductor to keep -the per-step launch count low; the generated Triton kernels match the GPU-side -cost of hand-written fused kernels once the model forward is long enough to -hide the compiled-function entry. - -Compile configuration (benchmarked against eager and other modes on H200): - -- ``mode="max-autotune-no-cudagraphs"``: CUDA graphs are deliberately - disabled -- the update mutates persistent per-slot state in place and the - gather's output is consumed outside the compiled region, both of which are - unsafe with cudagraph static buffers (outputs get overwritten by subsequent - replays). -- ``torch._dynamo.mark_dynamic`` on the batch-varying dimensions, so per-step - changes in batch composition do not trigger recompilation while the - remaining dimensions stay specialized. - -Compilation happens lazily on the first decay-active request (roughly a -second); non-decay workloads never trigger it. -""" - -import torch - - -@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) - - -@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) - - -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) - _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, - ) - - -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 _top_p_decay_gather_impl(runtime_top_p, is_decay_slot, static_top_p, slots) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py index 5307876ae61d..8cfcca1e13b0 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 3b1e5eda0078..105e61941379 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -88,7 +88,6 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests -from .ops.top_p_decay import top_p_decay_update from .sampling_utils import ( BEAM_SEARCH_PAD_TOKEN, GREEDY, @@ -4680,7 +4679,7 @@ def _update_top_p_decay_after_sample( if not self._top_p_decay_slots: return store = self.store.top_p_decay_store - top_p_decay_update( + Fusions.top_p_decay_update( runtime_top_p=store.runtime_top_p_decay_cuda, initial_top_p=store.initial_top_p_decay_cuda, top_p_decay=store.top_p_decay_cuda, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index b3a9f1bfe6c2..7ca58c7edc08 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -39,7 +39,6 @@ top_p_renorm_probs_op, top_p_sampling_from_probs_op, ) -from tensorrt_llm._torch.pyexecutor.sampler.ops.top_p_decay import top_p_decay_gather from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import ( GREEDY_TEMPERATURE_THRESHOLD, BeamSearchMetadata, @@ -411,7 +410,7 @@ def _maybe_apply_top_p_decay(self, group_metadata: Optional[StrategyMetadata]) - self._top_p.shape, group_metadata.slots.shape, ) - self._top_p = top_p_decay_gather( + self._top_p = Fusions.top_p_decay_gather( runtime_top_p=group_metadata.runtime_top_p, is_decay_slot=group_metadata.is_decay_slot, static_top_p=self._top_p, From 6133665574235f88448bfaaf9c5556163041cc48 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Jul 2026 00:33:33 -0700 Subject: [PATCH 8/9] [TRTLLM-13231][chore] Fix stale kernel wording and an unrelated whitespace change - TopPDecayStore lifecycle step 1 still said 'fused kernels'; the CUDA kernels were replaced by torch.compile'd ops. - Restore a blank line in setup_sampler_step accidentally dropped from the upstream code, keeping the diff against main minimal. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 105e61941379..5996b2a012bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2287,7 +2287,7 @@ class TopPDecayStore: cleared then re-set for the newly-admitted slots -- both the host-side ``TorchSampler._top_p_decay_slots`` set (an O(1) hot-path early-out) and its device mirror ``is_top_p_decay_slot_cuda`` (the gate the - fused kernels use, so the hot path needs no host-side filtering) -- + fused ops use, so the hot path needs no host-side filtering) -- and the per-slot buffers are initialized. This clear-then-set also covers slot reuse: stale entries from a prior occupant are never consumed. @@ -3092,6 +3092,7 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: seq_slots: list[int] = [] # Used for beam search updates max_prompt_len: int = 0 + # Prepare finish reasons handler self._finish_reasons_handler.setup_new_request_handling() for request in new_requests: From 49c0c79298bda43c36bc84690fdda76697281ab3 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Jul 2026 04:05:25 -0700 Subject: [PATCH 9/9] [TRTLLM-13231][fix] Cast top_p_decay group bounds to int for slice indexing mypy rejected slicing seq_slots_cuda with Tensor.item() results (Slice index must be an integer). Cast first_req/last_req to int, matching the existing int()/cast(int, ...) pattern used elsewhere in this file. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 5996b2a012bb..2f885dcbfffa 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -4312,8 +4312,8 @@ def _build_top_p_decay_metadata( # (group_req_indices being sorted ascending) a contiguous group's slots # are just a slice of seq_slots_cuda -- no host layout build and no H2D # copy. - first_req = group_req_indices[0].item() - last_req = group_req_indices[-1].item() + first_req = int(group_req_indices[0].item()) + last_req = int(group_req_indices[-1].item()) group_steps = req_num_steps[group_req_indices] if last_req - first_req + 1 == group_req_indices.size(0) and ( group_steps.max().item() == 1