From b7ee51970ec3ef49c6d5dfe05a50331ef76bea0d Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <zhaoyangw@nvidia.com>
Date: Mon, 13 Jul 2026 02:49:15 -0700
Subject: [PATCH 1/5] [TRTLLM-13212][refactor] Unify sampler ops: dedup
 wrappers, merge flashinfer op variants, drop dead trtllm op

Follow-ups promised in PR #15542 review:

- vanilla: fold top_k_/top_p_/temperature_sampling_batch thin wrappers into
  top_k_top_p_sampling_batch via default args (top_k=None disables top-k,
  top_p=1 disables top-p); drop the pass-through greedy() wrapper
  (greedy re-export now aliases greedy_search_sampling_batch directly).
- flashinfer ops: merge the seed/offset and generator variants of each
  sampling op into a single signature with keyword-only randomness args,
  matching flashinfer's own unified API (>= 0.6.4 accepts both; explicit
  seed/offset takes precedence). 6 wrappers -> 4.
- Fix seed/offset type hints on sampling_batch_spec_dec_one_model
  (Optional[int] -> Optional[torch.Tensor]; callers pass device tensors
  for CUDA graph compatibility).
- Remove dead trtllm::compute_probs_from_logits_op chain (no callers
  anywhere): Python register_fake, thop wrapper + TORCH_LIBRARY
  registration, and the exclusive kernel chain in dynamicTreeKernels.cu
  (topKProbStage1/2, invokeTopKTopPMaskingForProbs, applyTopKTopPForProbOp,
  computeSoftmaxForProbOp, computeProbsFromLogits) plus now-unused
  torch/ATen includes.
- Refresh stale docstring references in sanitize_top_k.

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
---
 .../speculativeDecoding/dynamicTreeKernels.cu | 476 ------------------
 cpp/tensorrt_llm/thop/dynamicTreeOp.cpp       |  40 --
 .../_torch/auto_deploy/shim/demollm.py        |   4 +-
 .../_torch/custom_ops/cpp_custom_ops.py       |   8 -
 .../pyexecutor/sampler/ops/flashinfer.py      | 136 ++---
 .../_torch/pyexecutor/sampler/ops/vanilla.py  |  68 +--
 .../pyexecutor/sampler/sampling_utils.py      |  63 ++-
 7 files changed, 114 insertions(+), 681 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu
index 3d167660e521..caa090f0e25c 100644
--- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu
+++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu
@@ -29,12 +29,10 @@
 #include "tensorrt_llm/common/reduceKernelUtils.cuh"
 #include "tensorrt_llm/common/vec_dtypes.cuh"
 #include "tensorrt_llm/kernels/decodingCommon.h"
-#include <ATen/cuda/CUDAContext.h>
 #include <algorithm>
 #include <cfloat>
 #include <cstdint>
 #include <limits>
-#include <torch/extension.h>
 TRTLLM_NAMESPACE_BEGIN
 
 using namespace tensorrt_llm::common;
@@ -43,480 +41,6 @@ using namespace tensorrt_llm::runtime;
 namespace kernels::speculative_decoding
 {
 
-// ---------------------------------------------------------------------------
-// Two-stage top-k / top-p masking kernels
-// Mirrors the approach in invokeBatchTopKSampling (samplingTopKKernels.cu),
-// but outputs a masked logits tensor instead of sampling a token.
-// ---------------------------------------------------------------------------
-
-// Stage 1: Parallel top-k reduction across BLOCKS_PER_BEAM_ blocks per row.
-// Each block handles (vocabSize / BLOCKS_PER_BEAM_) elements and finds its
-// local top-k, writing (global_index, logit_value) pairs into the tmp buffers.
-template <typename T, int32_t BLOCK_SIZE_, int32_t BLOCKS_PER_BEAM_>
-__global__ void topKProbStage1(T const* __restrict__ logits, T* tmpLogProbs, int32_t* topKTmpIdBuf, T* topKTmpValBuf,
-    int32_t maxTopK, int32_t const* topKs, int32_t vocabSize)
-{
-    typedef cub::BlockReduce<TopK_2<T>, BLOCK_SIZE_> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage tempStorage;
-
-    auto const tid = static_cast<int32_t>(threadIdx.x);
-    auto const bid = static_cast<int32_t>(blockIdx.x);
-    auto const rowId = bid / BLOCKS_PER_BEAM_;
-    auto const blockLane = bid % BLOCKS_PER_BEAM_; // chunk index within the row
-
-    auto const k = (topKs != nullptr) ? topKs[rowId] : maxTopK;
-
-    bool const IS_FP16 = std::is_same<T, half>::value;
-    T const MAX_T_VAL = IS_FP16 ? HALF_FLT_MAX : FLT_MAX;
-
-    // Base offset into the flat (nRows * vocabSize) logits array for this row.
-    auto const rowOffset = rowId * vocabSize;
-    // Base offset into the tmp buffers for this (row, blockLane).
-    auto const tmpIdxBase = rowId * BLOCKS_PER_BEAM_ * maxTopK + blockLane * k;
-
-    // Copy this block's chunk of logits into tmpLogProbs scratch space.
-    for (auto elemId = tid + blockLane * BLOCK_SIZE_; elemId < vocabSize; elemId += BLOCK_SIZE_ * BLOCKS_PER_BEAM_)
-    {
-        tmpLogProbs[rowOffset + elemId] = logits[rowOffset + elemId];
-    }
-    __syncthreads();
-
-    // Iteratively find the top-k values via max-reduction, zeroing each found max.
-    TopK_2<T> partial;
-    for (int32_t ite = 0; ite < k; ite++)
-    {
-        partial.init();
-        for (auto elemId = tid + blockLane * BLOCK_SIZE_; elemId < vocabSize; elemId += BLOCK_SIZE_ * BLOCKS_PER_BEAM_)
-        {
-            partial.insert(tmpLogProbs[rowOffset + elemId], rowOffset + elemId);
-        }
-
-        TopK_2<T> total = BlockReduce(tempStorage).Reduce(partial, reduce_topk_op_2<T>);
-
-        if (tid == 0)
-        {
-            topKTmpIdBuf[tmpIdxBase + ite] = total.p;  // global index (rowOffset + vocabIdx)
-            topKTmpValBuf[tmpIdxBase + ite] = total.u; // logit value
-            if (total.p >= 0)
-            {
-                tmpLogProbs[total.p] = -MAX_T_VAL; // zero out so next iteration finds next-best
-            }
-        }
-        __syncthreads();
-    }
-}
-
-// Stage 2: Merge BLOCKS_PER_BEAM_ * k candidates per row, apply optional top-p,
-// then scatter selected logit values back to an output logits tensor (all other
-// positions are set to -inf so that a subsequent softmax produces 0 probability).
-template <typename T, int32_t BLOCK_SIZE_, int32_t BLOCKS_PER_BEAM_>
-__global__ void topKProbStage2ForLogits(int32_t const* __restrict__ topKTmpIdBuf, T* topKTmpValBuf, float* outputLogits,
-    int32_t maxTopK, int32_t const* topKs, float const* topPs, int32_t vocabSize)
-{
-    bool const IS_FP16 = std::is_same<T, half>::value;
-    T const MAX_T_VAL = IS_FP16 ? HALF_FLT_MAX : FLT_MAX;
-
-    auto const tid = static_cast<int32_t>(threadIdx.x);
-    auto const rowId = static_cast<int32_t>(blockIdx.x);
-
-    auto const k = (topKs != nullptr) ? topKs[rowId] : maxTopK;
-    // size: number of valid candidates written by Stage 1 for this row.
-    // stride: row pitch in the tmp buffers (same as in invokeBatchTopKSampling).
-    auto const size = k * BLOCKS_PER_BEAM_;
-    auto const stride = maxTopK * BLOCKS_PER_BEAM_;
-
-    typedef cub::BlockReduce<TopK_2<float>, BLOCK_SIZE_> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage tempStorage;
-    extern __shared__ char sharedArray[];
-    // Shared layout: sId[maxTopK] | sVal2[maxTopK]
-    auto* sId = reinterpret_cast<int32_t*>(sharedArray);
-    auto* sVal2 = reinterpret_cast<float*>(sId + maxTopK);
-
-    // Pointer to this row's candidates in the tmp value buffer (modified in-place during reduction).
-    T* sVal = topKTmpValBuf + rowId * stride;
-
-    // Step 1: Initialize output row to -inf (all threads cooperate for bandwidth).
-    float* outRow = outputLogits + rowId * vocabSize;
-    float const negInf = -std::numeric_limits<float>::infinity();
-    for (int32_t i = tid; i < vocabSize; i += BLOCK_SIZE_)
-    {
-        outRow[i] = negInf;
-    }
-    __syncthreads();
-
-    // Step 2: k-round block-reduction over the k * BLOCKS_PER_BEAM_ valid candidates.
-    // (Only the first 'size' entries of the row's tmp buffer were written by Stage 1.)
-    TopK_2<float> partial;
-    __shared__ float sMaxLogit;
-    for (int32_t ite = 0; ite < k; ite++)
-    {
-        partial.init();
-        for (int32_t i = tid; i < size; i += BLOCK_SIZE_)
-        {
-            partial.insert(static_cast<float>(sVal[i]), i);
-        }
-
-        TopK_2<float> total = BlockReduce(tempStorage).Reduce(partial, reduce_topk_op_2<float>);
-
-        if (tid == 0)
-        {
-            if (ite == 0)
-            {
-                sMaxLogit = total.u;
-            }
-            sId[ite] = total.p;
-            sVal[total.p] = -MAX_T_VAL; // zero out so next iteration finds next-best
-            sVal2[ite] = total.u;       // store raw logit value (not exponentiated)
-        }
-        __syncthreads();
-    }
-
-    // Step 3: Determine top-p cutoff (tid=0 only).
-    // sVal2 contains logit values in descending order; we exponentiate to get unnormalized probs.
-    if (tid == 0)
-    {
-        int32_t cutoff = k;
-        if (topPs != nullptr)
-        {
-            float const topP = topPs[rowId];
-            if (topP < 1.0f)
-            {
-                // Compute unnormalized probabilities and their sum.
-                float sSum = 0.0f;
-                for (int32_t ki = 0; ki < k; ki++)
-                {
-                    sVal2[ki] = __expf(sVal2[ki] - sMaxLogit); // reuse sVal2 to hold exp probs
-                    sSum += sVal2[ki];
-                }
-                // Walk in descending-probability order; stop as soon as cumulative prob >= topP.
-                float cumProb = 0.0f;
-                for (int32_t ki = 0; ki < k; ki++)
-                {
-                    cumProb += sVal2[ki] / sSum;
-                    if (cumProb >= topP)
-                    {
-                        cutoff = ki + 1; // always keep at least this token
-                        break;
-                    }
-                }
-            }
-        }
-
-        // Step 4: Scatter selected logit values back to output.
-        // topKTmpIdBuf stores (rowOffset + vocabIdx); recover vocabIdx with % vocabSize.
-        auto const rowStride = rowId * stride;
-        for (int32_t ki = 0; ki < cutoff; ki++)
-        {
-            auto const candidateIdx = sId[ki];
-            auto const globalIdx = topKTmpIdBuf[rowStride + candidateIdx];
-            if (globalIdx >= 0)
-            {
-                auto const vocabIdx = globalIdx % vocabSize;
-                // sVal2 was overwritten with exp probs when topP < 1; we need the original logit.
-                // Re-read from the original tmp buffer — the stored value IS the logit (set in Stage 1).
-                // However sVal[candidateIdx] was zeroed during Stage 2 reduction; but
-                // topKTmpValBuf still holds the original value at that index (sVal points there).
-                // We stored the logit as sVal2[ite] = total.u BEFORE any exp, so if topP was not
-                // applied we can use sVal2[ki] directly. If topP was applied, sVal2[ki] now holds
-                // the exp prob — we cannot recover the logit. To handle both cases cleanly we use
-                // log(sVal2[ki]) + sMaxLogit when topPs was applied, otherwise sVal2[ki] directly.
-                float logitVal;
-                if (topPs != nullptr && topPs[rowId] < 1.0f)
-                {
-                    // sVal2[ki] = exp(logit - sMaxLogit), so logit = log(sVal2[ki]) + sMaxLogit
-                    logitVal = __logf(sVal2[ki]) + sMaxLogit;
-                }
-                else
-                {
-                    logitVal = sVal2[ki]; // still the raw logit
-                }
-                outRow[vocabIdx] = logitVal;
-            }
-        }
-    }
-}
-
-#define CASE_K_PROB(K_MAX, BLOCK_SIZE_1_, BLOCK_SIZE_2_, BLOCKS_PER_BEAM_)                                             \
-    do                                                                                                                 \
-    {                                                                                                                  \
-        topKProbStage1<T, BLOCK_SIZE_1_, BLOCKS_PER_BEAM_>                                                             \
-            <<<dim3(nRows* BLOCKS_PER_BEAM_, 1), BLOCK_SIZE_1_, 0, stream>>>(                                          \
-                logits, tmpLogProbs, topKTmpIdBuf, topKTmpValBuf, maxTopK, topKs, vocabSize);                          \
-        topKProbStage2ForLogits<T, BLOCK_SIZE_2_, BLOCKS_PER_BEAM_>                                                    \
-            <<<dim3(nRows, 1), BLOCK_SIZE_2_, K_MAX * (sizeof(int32_t) + sizeof(float)), stream>>>(                    \
-                topKTmpIdBuf, topKTmpValBuf, outputLogits, maxTopK, topKs, topPs, vocabSize);                          \
-    } while (0)
-
-// Host launcher: allocates workspace tensors internally and dispatches the two-stage kernels.
-// logits        [nRows, vocabSize] – temperature-scaled input (float or half)
-// outputLogits  [nRows, vocabSize] – output: -inf everywhere except selected top-k-p positions
-// topKs         [nRows]            – per-row k values (int32, on device)
-// topPs         [nRows] or nullptr – per-row p values (float, on device)
-// maxTopK                          – maximum k across all rows (CPU scalar, 1–1024)
-template <typename T>
-void invokeTopKTopPMaskingForProbs(T const* logits, float* outputLogits, int32_t const* topKs, float const* topPs,
-    int32_t maxTopK, int32_t nRows, int32_t vocabSize, cudaStream_t stream)
-{
-    constexpr int32_t BLOCKS_PER_BEAM = 8;
-
-    // Workspace buffers (allocated as CUDA device tensors via ATen).
-    auto opts = at::TensorOptions().dtype(torch::kFloat32).device(at::kCUDA);
-    auto tmpLogProbsTensor = torch::empty({nRows * vocabSize}, opts);
-    auto topKTmpIdBufTensor
-        = torch::empty({nRows * BLOCKS_PER_BEAM * maxTopK}, at::TensorOptions().dtype(torch::kInt32).device(at::kCUDA));
-    // topKTmpValBuf uses the same dtype as T; we allocate as float and reinterpret for half if needed.
-    auto topKTmpValBufTensor = torch::empty({nRows * BLOCKS_PER_BEAM * maxTopK}, opts);
-
-    T* tmpLogProbs = reinterpret_cast<T*>(tmpLogProbsTensor.data_ptr<float>());
-    int32_t* topKTmpIdBuf = topKTmpIdBufTensor.data_ptr<int32_t>();
-    T* topKTmpValBuf = reinterpret_cast<T*>(topKTmpValBufTensor.data_ptr<float>());
-
-    int32_t logMaxTopK = 0;
-    int32_t recursor = maxTopK - 1;
-    while (recursor >>= 1)
-    {
-        ++logMaxTopK;
-    }
-
-    switch (logMaxTopK)
-    {
-    case 0:
-    case 1:
-    case 2:
-    case 3: // 0 < maxTopK <= 16
-        CASE_K_PROB(16, 128, 128, 8);
-        break;
-    case 4: // 16 < maxTopK <= 32
-        CASE_K_PROB(32, 256, 128, 8);
-        break;
-    case 5: // 32 < maxTopK <= 64
-        CASE_K_PROB(64, 256, 256, 8);
-        break;
-    case 6:
-    case 7:
-    case 8:
-    case 9: // 64 < maxTopK <= 1024
-        CASE_K_PROB(1024, 256, 256, 8);
-        break;
-    default: TLLM_CHECK_WITH_INFO(false, "topKProbMasking supports 1 <= k <= 1024 but got k=%d", maxTopK);
-    }
-}
-
-#undef CASE_K_PROB
-
-namespace
-{
-constexpr double kGreedyTempThreshold = 1e-4;
-
-torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits)
-{
-    TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
-    TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor");
-
-    auto probs = logits.contiguous().to(torch::kFloat32);
-    auto stream = at::cuda::getCurrentCUDAStream(probs.device().index());
-
-    BiasSoftmaxParams<float> biasSoftmaxParams;
-    biasSoftmaxParams.logits = probs.data_ptr<float>();
-    biasSoftmaxParams.probs = probs.data_ptr<float>();
-    biasSoftmaxParams.batchSize = static_cast<SizeType32>(probs.size(0));
-    biasSoftmaxParams.maxBatchSize = static_cast<SizeType32>(probs.size(0));
-    biasSoftmaxParams.maxBeamWidth = 1;
-    biasSoftmaxParams.vocabSize = static_cast<SizeType32>(probs.size(1));
-    biasSoftmaxParams.vocabSizePadded = static_cast<SizeType32>(probs.size(1));
-    biasSoftmaxParams.skipSoftMax = false;
-    biasSoftmaxParams.batchSlotsLogits = false;
-    biasSoftmaxParams.checkParams();
-
-    invokeAddBiasSoftMax(biasSoftmaxParams, stream);
-    return probs;
-}
-
-// Fast path for top-K (and optional top-P) filtering using torch::topk instead of a
-// full vocab-size sort.  kMax must be provided as a CPU integer (the caller computes it
-// via topK.max().item() on the Python side).  When kMax == 0 or kMax >= vocabSize the
-// function falls back to the original sort-based path.
-//
-// Key advantages over the full-sort path:
-//   1. torch::topk with small kMax is O(V * log kMax) vs O(V * log V) for full sort.
-//   2. The topk index tensor is [nRows, kMax] instead of [nRows, V] — much smaller.
-//   3. No scatter-back of sorted indices needed; masking is done directly on logits.
-//   4. For combined top-K + top-P, softmax/cumsum are computed on kMax values (not V).
-torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optional<torch::Tensor> const& topK,
-    torch::optional<torch::Tensor> const& topP, int32_t kMax)
-{
-    int64_t const vocabSize = logits.size(1);
-    // Host-only checks: the caller is expected to pass nullopt when filtering is fully
-    // disabled (see SpecMetadata.skip_top_k / skip_top_p). Probing the tensor contents
-    // via `.item<bool>()` here would force a host-device sync and break CUDA graph
-    // capture; the per-row `effectiveTopK` formula below already handles disabled rows.
-    bool const hasTopK = topK.has_value() && topK->defined();
-    bool const hasTopP = topP.has_value() && topP->defined();
-
-    if (!hasTopK && !hasTopP)
-    {
-        return logits;
-    }
-
-    torch::Tensor effectiveTopK;
-    if (hasTopK)
-    {
-        auto topKLong = topK->to(torch::kLong);
-        effectiveTopK
-            = torch::where(topKLong > 0, topKLong, torch::full_like(topKLong, vocabSize)).clamp_max(vocabSize);
-    }
-
-    // Fast path uses `topk(kMax)` which is unsafe when any row has effective top-k > kMax
-    // (i.e. disabled rows expand to the full vocab). Detecting this requires a tensor
-    // reduction + `.item<bool>()`, which is incompatible with CUDA graph capture. Only
-    // probe when the caller explicitly opted into the fast path via kMax > 0 (today only
-    // the dynamic-tree caller, which is not graph-captured).
-    bool hasDisabledTopKRows = false;
-    if (hasTopK && kMax > 0 && kMax < vocabSize)
-    {
-        auto topKLong = topK->to(torch::kLong);
-        hasDisabledTopKRows = topKLong.le(0).any().item<bool>();
-    }
-
-    if (hasTopK && !hasDisabledTopKRows && kMax > 0 && kMax < vocabSize)
-    {
-        // Fast topk path ─────────────────────────────────────────────────────────────
-        // topKValues/topKIdx: [nRows, kMax], values in descending order
-        auto [topKValues, topKIdx] = logits.topk(kMax, /*dim=*/-1, /*largest=*/true, /*sorted=*/true);
-
-        // validTopK[i, j]: True when position j falls within top-K[i] for row i
-        auto kArange = torch::arange(kMax, torch::TensorOptions().dtype(torch::kInt64).device(logits.device()))
-                           .unsqueeze(0);                          // [1, kMax]
-        auto kVals = effectiveTopK.to(torch::kInt64).unsqueeze(1); // [nRows, 1]
-        auto validTopK = kArange < kVals;                          // [nRows, kMax]
-
-        // Start with everything masked; scatter will unmark the kept positions.
-        auto mask = torch::ones(
-            {logits.size(0), vocabSize}, torch::TensorOptions().dtype(torch::kBool).device(logits.device()));
-
-        if (hasTopP)
-        {
-            // Compute top-P on the kMax descending-sorted values only (much cheaper).
-            // Positions beyond K[i] are treated as -inf so their probability ≈ 0.
-            auto validTopKValues = topKValues.masked_fill(~validTopK, -std::numeric_limits<float>::infinity());
-            auto sortedProbs = validTopKValues.softmax(/*dim=*/-1); // [nRows, kMax]
-            auto cumsum = sortedProbs.cumsum(/*dim=*/-1);           // [nRows, kMax]
-            // Mask positions where the cumulative probability *before* this token
-            // already reaches topP — i.e. we have enough probability mass already.
-            auto topPMask = (cumsum - sortedProbs) >= topP->unsqueeze(1); // [nRows, kMax]
-            topPMask.select(/*dim=*/1, /*index=*/0).fill_(false);         // always keep the top-1 token
-            // combinedMask: True  → mask this vocab position
-            //               False → keep this vocab position
-            auto combinedMask = topPMask | (~validTopK); // [nRows, kMax]
-            mask.scatter_(/*dim=*/1, /*index=*/topKIdx, /*src=*/combinedMask);
-        }
-        else
-        {
-            // Top-K only: unmark the first K[i] positions (those within validTopK).
-            // ~validTopK is True for positions j >= K[i] → they should stay masked.
-            mask.scatter_(/*dim=*/1, /*index=*/topKIdx, /*src=*/(~validTopK));
-        }
-
-        return logits.masked_fill(mask, -std::numeric_limits<float>::infinity());
-    }
-
-    // Fallback: full-sort path (used for top-P only, or when kMax == 0) ────────────
-    auto sortResult = logits.sort(/*dim=*/-1, /*descending=*/false);
-    auto logitsSort = std::get<0>(sortResult);
-    auto logitsIdx = std::get<1>(sortResult);
-
-    if (hasTopK)
-    {
-        auto topKMask = logitsSort.size(1) - effectiveTopK;
-        topKMask = topKMask.clamp_min(0);
-        auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1));
-        auto mask = logitsSort < topKThreshold;
-        logitsSort.masked_fill_(mask, -std::numeric_limits<float>::infinity());
-    }
-
-    if (hasTopP)
-    {
-        auto probsSort = logitsSort.softmax(/*dim=*/-1);
-        auto probsSum = probsSort.cumsum(/*dim=*/-1, /*dtype=*/probsSort.scalar_type());
-        auto topPMask = probsSum <= (1.0 - topP->unsqueeze(1));
-        topPMask.select(/*dim=*/1, /*index=*/logitsSort.size(1) - 1).fill_(false);
-        logitsSort.masked_fill_(topPMask, -std::numeric_limits<float>::infinity());
-    }
-
-    return logitsSort.scatter(/*dim=*/-1, /*index=*/logitsIdx, /*src=*/logitsSort);
-}
-
-} // namespace
-
-torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor const& temperatures,
-    torch::optional<torch::Tensor> const& topK, torch::optional<torch::Tensor> const& topP, bool skipTemperature,
-    int32_t kMax)
-{
-    TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
-    TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor");
-    TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor");
-    TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor");
-    TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch");
-    if (topK.has_value() && topK->defined())
-    {
-        TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor");
-        TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor");
-        TORCH_CHECK(topK->size(0) == logits.size(0), "top_k and logits size mismatch");
-    }
-    if (topP.has_value() && topP->defined())
-    {
-        TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor");
-        TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor");
-        TORCH_CHECK(topP->size(0) == logits.size(0), "top_p and logits size mismatch");
-    }
-
-    auto const isGreedy = temperatures <= kGreedyTempThreshold;
-    auto const safeTemperatures = torch::where(isGreedy, torch::ones_like(temperatures), temperatures);
-    auto scaledLogits
-        = (skipTemperature ? logits : logits.div(safeTemperatures.unsqueeze(1))).contiguous().to(torch::kFloat32);
-
-    int64_t const vocabSize = scaledLogits.size(1);
-    int64_t const nRows = scaledLogits.size(0);
-    // Host-only presence checks; see comment in applyTopKTopPForProbOp() for why we
-    // avoid probing tensor contents (would sync and break CUDA graph capture).
-    bool const hasTopKPresence = topK.has_value() && topK->defined();
-    bool const hasTopPPresence = topP.has_value() && topP->defined();
-
-    // The kernel path produces -inf for rows whose top_k value is 0, so it is only
-    // safe when every row has an active top_k filter. Determining that requires a
-    // host-device sync, so only probe when the caller has opted into the kernel
-    // path (kMax > 0). The kMax > 0 callers (dynamic-tree) are not graph-captured.
-    bool useKernelPath = false;
-    if (hasTopKPresence && kMax > 0 && kMax < vocabSize)
-    {
-        useKernelPath = torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item<bool>();
-    }
-
-    torch::Tensor maskedLogits;
-    if (useKernelPath)
-    {
-        // Two-stage CUDA top-k/top-p masking (mirrors invokeBatchTopKSampling).
-        maskedLogits = torch::empty_like(scaledLogits);
-        auto topKForKernel = topK->to(torch::kInt32).contiguous();
-        auto topPForKernel = hasTopPPresence ? topP->to(torch::kFloat32).contiguous() : torch::Tensor();
-        auto stream = at::cuda::getCurrentCUDAStream(scaledLogits.device().index());
-        invokeTopKTopPMaskingForProbs<float>(scaledLogits.data_ptr<float>(), maskedLogits.data_ptr<float>(),
-            topKForKernel.data_ptr<int32_t>(), hasTopPPresence ? topPForKernel.data_ptr<float>() : nullptr, kMax,
-            static_cast<int32_t>(nRows), static_cast<int32_t>(vocabSize), stream);
-    }
-    else
-    {
-        // Fallback: PyTorch-based sort path (top-P only or kMax == 0).
-        maskedLogits = applyTopKTopPForProbOp(scaledLogits, topK, topP, kMax);
-    }
-
-    auto probs = computeSoftmaxForProbOp(maskedLogits);
-
-    auto argmaxIds = maskedLogits.argmax(/*dim=*/-1, /*keepdim=*/true);
-    auto oneHot = torch::zeros_like(probs).scatter_(1, argmaxIds, 1.0);
-    return torch::where(isGreedy.unsqueeze(1), oneHot, probs);
-}
-
 //! \param parentList           [in]  layer-wise parent indices [bs, topK*(depth-1)+1]
 //! \param selectedIndex        [in]  resampled history buffer indices [bs, draftTokenNum-1]
 //! \param treeMask             [out] attention mask (which nodes each node can see)
diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp
index 0b862108a4e6..0a236ced8c35 100644
--- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp
+++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp
@@ -25,13 +25,6 @@ namespace tk = tensorrt_llm::kernels::speculative_decoding;
 
 TRTLLM_NAMESPACE_BEGIN
 
-namespace kernels::speculative_decoding
-{
-th::Tensor computeProbsFromLogits(th::Tensor const& logits, th::Tensor const& temperatures,
-    th::optional<th::Tensor> const& topK, th::optional<th::Tensor> const& topP, bool skipTemperature,
-    runtime::SizeType32 kMax);
-} // namespace kernels::speculative_decoding
-
 namespace torch_ext
 {
 
@@ -126,26 +119,6 @@ void verify_dynamic_tree_greedy_out_packed_op(th::Tensor& candidates, th::Tensor
         targetPredict.data_ptr<int32_t>(), treeValid.data_ptr<bool>(), batchSize, numDraftTokens, numSpecStep, stream);
 }
 
-th::Tensor compute_probs_from_logits_op(th::Tensor logits, th::Tensor temperatures, th::optional<th::Tensor> topK,
-    th::optional<th::Tensor> topP, bool skipTemperature)
-{
-    TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
-    TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor");
-    TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor");
-    TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor");
-    TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch");
-    if (topK.has_value() && topK->defined())
-    {
-        TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor");
-    }
-    if (topP.has_value() && topP->defined())
-    {
-        TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor");
-    }
-
-    return tk::computeProbsFromLogits(logits, temperatures, topK, topP, skipTemperature, /*kMax=*/0);
-}
-
 //! \brief Target-only rejection sampling verify op (no draft probabilities needed).
 void verify_dynamic_tree_rejection_out_op(th::Tensor& draftTokens, th::Tensor& targetProbs,
     th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, th::Tensor& treeValid, th::Tensor& acceptIndex,
@@ -286,16 +259,3 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("verify_dynamic_tree_rejection_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_rejection_out_op);
 }
-
-TORCH_LIBRARY_FRAGMENT(trtllm, m)
-{
-    m.def(
-        "compute_probs_from_logits_op("
-        "Tensor logits, Tensor temperatures, Tensor? top_k=None, Tensor? top_p=None, "
-        "bool skip_temperature=False) -> Tensor");
-}
-
-TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
-{
-    m.impl("compute_probs_from_logits_op", &tensorrt_llm::torch_ext::compute_probs_from_logits_op);
-}
diff --git a/tensorrt_llm/_torch/auto_deploy/shim/demollm.py b/tensorrt_llm/_torch/auto_deploy/shim/demollm.py
index 9f12f444b880..104ca0c7b266 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/demollm.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/demollm.py
@@ -24,7 +24,7 @@
 
 from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import (
     greedy_search_sampling_batch,
-    top_k_sampling_batch,
+    top_k_top_p_sampling_batch,
 )
 from tensorrt_llm.executor import GenerationExecutor
 from tensorrt_llm.executor.request import GenerationRequest
@@ -312,7 +312,7 @@ def _sample(
         logits_shape = logits.shape
         logits = logits.view(-1, logits_shape[-1])  # sampling_batch expects 2D logits
         if isinstance(sampling_params.top_k, int) and sampling_params.top_k > 1:
-            idx_next, probs = top_k_sampling_batch(
+            idx_next, probs = top_k_top_p_sampling_batch(
                 logits, top_k=sampling_params.top_k, temperature=1.0
             )
         else:
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index e82d81e219b2..7e1044e460c6 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1388,11 +1388,3 @@ def _(like: torch.Tensor,
         out_shape = shape if shape is not None else list(like.shape)
         dtype = out_dtype if out_dtype is not None else like.dtype
         return like.new_empty(out_shape, dtype=dtype), output_buffer_kind
-
-    @torch.library.register_fake("trtllm::compute_probs_from_logits_op")
-    def _(logits: torch.Tensor,
-          temperatures: torch.Tensor,
-          top_k: Optional[torch.Tensor] = None,
-          top_p: Optional[torch.Tensor] = None,
-          skip_temperature: bool = False) -> torch.Tensor:
-        return logits.new_empty(list(logits.shape), dtype=torch.float32)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
index 902472d8f562..4d8226bc4673 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
@@ -16,9 +16,18 @@
 
 These ops depend on flashinfer; the import is guarded so the module stays
 importable without it.
+
+Randomness can be supplied either way (flashinfer accepts both in one
+signature; explicit ``seed``/``offset`` take precedence over ``generator``):
+
+- ``generator``: stateful host-side ``torch.Generator``, for eager paths.
+- ``seed``/``offset``: stateless device tensors, required under CUDA graph
+  capture (a ``torch.Generator`` advances host-side at launch time, so its
+  state would be frozen into the graph and every replay would reuse the same
+  random values).
 """
 
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
@@ -27,91 +36,59 @@
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer.sampling
 
+SeedOrTensor = Union[int, torch.Tensor]
+
 
 def top_k_top_p_sampling_from_logits_op(
     logits: torch.Tensor,
     top_k: torch.Tensor,
     top_p: torch.Tensor,
-    seed: Optional[int] = None,
-    offset: Optional[int] = None,
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
+    check_nan: bool = False,
 ) -> torch.Tensor:
     tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
-        logits, top_k, top_p, seed=seed, offset=offset
+        logits,
+        top_k=top_k,
+        top_p=top_p,
+        filter_apply_order="top_k_first",
+        deterministic=True,
+        check_nan=check_nan,
+        generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
 def sampling_from_probs_op(
     probs: torch.Tensor,
-    seed: Optional[torch.Tensor] = None,
-    offset: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
-        probs, deterministic=True, seed=seed, offset=offset
-    )
-    return tokens
-
-
-def softmax_op(
-    logits: torch.Tensor,
-    temperature: Optional[torch.Tensor],
-) -> torch.Tensor:
-    probs: torch.Tensor = flashinfer.sampling.softmax(
-        logits, temperature, enable_pdl=get_env_enable_pdl()
-    )
-    return probs
-
-
-def top_k_mask_logits_op(
-    logits: torch.Tensor,
-    top_k: torch.Tensor,
-) -> torch.Tensor:
-    masked: torch.Tensor = flashinfer.sampling.top_k_mask_logits(logits, top_k)
-    return masked
-
-
-def top_p_renorm_probs_op(
-    probs: torch.Tensor,
-    top_p: torch.Tensor,
-) -> torch.Tensor:
-    renormed: torch.Tensor = flashinfer.sampling.top_p_renorm_probs(probs, top_p)
-    return renormed
-
-
-def sampling_from_probs_generator_op(
-    probs: torch.Tensor,
-    generator: Optional[torch.Generator],
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
     tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
-        probs, deterministic=True, generator=generator, check_nan=check_nan
-    )
-    return tokens
-
-
-def top_k_top_p_sampling_from_logits_with_generator_op(
-    logits: torch.Tensor,
-    top_k: torch.Tensor,
-    top_p: torch.Tensor,
-    generator: Optional[torch.Generator],
-    check_nan: bool = False,
-) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
-        logits,
-        top_k=top_k,
-        top_p=top_p,
-        filter_apply_order="top_k_first",
+        probs,
         deterministic=True,
         check_nan=check_nan,
         generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
-def top_k_sampling_from_probs_generator_op(
+def top_k_sampling_from_probs_op(
     probs: torch.Tensor,
     top_k: torch.Tensor,
-    generator: Optional[torch.Generator],
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
     tokens: torch.Tensor = flashinfer.sampling.top_k_sampling_from_probs(
@@ -120,14 +97,19 @@ def top_k_sampling_from_probs_generator_op(
         deterministic=True,
         check_nan=check_nan,
         generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
-def top_p_sampling_from_probs_generator_op(
+def top_p_sampling_from_probs_op(
     probs: torch.Tensor,
     top_p: torch.Tensor,
-    generator: Optional[torch.Generator],
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
     tokens: torch.Tensor = flashinfer.sampling.top_p_sampling_from_probs(
@@ -136,10 +118,38 @@ def top_p_sampling_from_probs_generator_op(
         deterministic=True,
         check_nan=check_nan,
         generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
+def softmax_op(
+    logits: torch.Tensor,
+    temperature: Optional[torch.Tensor],
+) -> torch.Tensor:
+    probs: torch.Tensor = flashinfer.sampling.softmax(
+        logits, temperature, enable_pdl=get_env_enable_pdl()
+    )
+    return probs
+
+
+def top_k_mask_logits_op(
+    logits: torch.Tensor,
+    top_k: torch.Tensor,
+) -> torch.Tensor:
+    masked: torch.Tensor = flashinfer.sampling.top_k_mask_logits(logits, top_k)
+    return masked
+
+
+def top_p_renorm_probs_op(
+    probs: torch.Tensor,
+    top_p: torch.Tensor,
+) -> torch.Tensor:
+    renormed: torch.Tensor = flashinfer.sampling.top_p_renorm_probs(probs, top_p)
+    return renormed
+
+
 def compute_probs_from_logits_op(
     logits: torch.Tensor,
     temperatures: torch.Tensor,
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
index ffb962a60fc6..1232ed37ea4e 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
@@ -50,66 +50,27 @@ class BeamSearchMetadata(StrategyMetadata):
     beam_idx_arange: torch.Tensor
 
 
-def top_k_sampling_batch(
-    logits: torch.Tensor,
-    *,
-    top_k: int,
-    temperature: float,
-    generator: Optional[torch.Generator] = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return top_k_top_p_sampling_batch(
-        logits,
-        top_k=top_k,
-        temperature=temperature,
-        generator=generator,
-        top_p=1,
-    )
-
-
-def top_p_sampling_batch(
-    logits: torch.Tensor,
-    *,
-    top_p: float,
-    temperature: float,
-    generator: Optional[torch.Generator] = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return top_k_top_p_sampling_batch(
-        logits,
-        top_p=top_p,
-        top_k=logits.size(1),
-        temperature=temperature,
-        generator=generator,
-    )
-
-
-def temperature_sampling_batch(
-    logits: torch.Tensor,
-    *,
-    temperature: float,
-    generator: Optional[torch.Generator] = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return top_k_top_p_sampling_batch(
-        logits,
-        top_p=1,
-        top_k=logits.size(1),
-        temperature=temperature,
-        generator=generator,
-    )
-
-
 def top_k_top_p_sampling_batch(
     logits: torch.Tensor,
     *,
-    top_k: int,
-    top_p: float,
     temperature: float,
+    top_k: Optional[int] = None,
+    top_p: float = 1.0,
     generator: Optional[torch.Generator] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Temperature + optional top-k / top-p filtering + multinomial sampling.
+
+    ``top_k=None`` (or ``vocab_size``) disables top-k filtering; ``top_p=1``
+    disables top-p filtering. With both disabled this is plain temperature
+    sampling.
+    """
     logits_dim = logits.dim()
     assert logits_dim == 2, "logits should be 2D: [batch_size, vocab_size]"
     assert temperature > 0, "non-greedy sampling requires valid temperature"
     logits = logits / max(temperature, 1e-5)
     batch_size, vocab_size = logits.size()
+    if top_k is None:
+        top_k = vocab_size
 
     assert top_k > 1, "non-greedy sampling requires valid top_k"
     need_top_k = top_k < vocab_size
@@ -325,15 +286,6 @@ def sample_rejected(
 _GREEDY_TEMPERATURE_THRESHOLD = 1e-4
 
 
-def greedy(
-    logits: torch.Tensor,
-    *,
-    return_probs: bool = True,
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-    """Greedy decoding; returns exact one-hot when return_probs=True."""
-    return greedy_search_sampling_batch(logits, return_probs=return_probs)
-
-
 def _safely_apply_temperature_inplace(
     logits_inout: torch.Tensor, temp: torch.Tensor
 ) -> torch.Tensor:
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
index 0c4d59b0ecc4..f27e9bb33b7b 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
@@ -31,23 +31,23 @@
 # These op wrappers are safe to import without flashinfer installed; they are
 # only called on the flashinfer sampler / speculative-worker paths.
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
-    sampling_from_probs_generator_op as sampling_from_probs_generator_op,
+    sampling_from_probs_op as sampling_from_probs_op,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import softmax_op as softmax_op
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
     top_k_mask_logits_op as top_k_mask_logits_op,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
-    top_k_sampling_from_probs_generator_op as top_k_sampling_from_probs_generator_op,
+    top_k_sampling_from_probs_op as top_k_sampling_from_probs_op,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
-    top_k_top_p_sampling_from_logits_with_generator_op as top_k_top_p_sampling_from_logits_with_generator_op,  # noqa: E501
+    top_k_top_p_sampling_from_logits_op as top_k_top_p_sampling_from_logits_op,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
     top_p_renorm_probs_op as top_p_renorm_probs_op,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
-    top_p_sampling_from_probs_generator_op as top_p_sampling_from_probs_generator_op,
+    top_p_sampling_from_probs_op as top_p_sampling_from_probs_op,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
     BeamSearchMetadata as BeamSearchMetadata,
@@ -60,23 +60,13 @@
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
     get_rejected_indices as get_rejected_indices,
 )
-from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import greedy as _torch_greedy
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
     greedy_search_sampling_batch as greedy_search_sampling_batch,
 )
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import sample_rejected as sample_rejected
-from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
-    temperature_sampling_batch as temperature_sampling_batch,
-)
-from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
-    top_k_sampling_batch as top_k_sampling_batch,
-)
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
     top_k_top_p_sampling_batch as top_k_top_p_sampling_batch,
 )
-from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
-    top_p_sampling_batch as top_p_sampling_batch,
-)
 from tensorrt_llm._utils import prefer_pinned
 from tensorrt_llm.sampling_params import SamplingParams
 
@@ -179,14 +169,14 @@ def sample(
     # 'cast' needed b/c of https://github.com/python/mypy/issues/19081
     match strategy:
         case ("top_k", top_k, temperature):
-            tokens, softmax = top_k_sampling_batch(
+            tokens, softmax = top_k_top_p_sampling_batch(
                 logits,
                 top_k=cast(int, top_k),
                 temperature=cast(float, temperature),
                 generator=generator,
             )
         case ("top_p", top_p, temperature):
-            tokens, softmax = top_p_sampling_batch(
+            tokens, softmax = top_k_top_p_sampling_batch(
                 logits,
                 top_p=cast(float, top_p),
                 generator=generator,
@@ -201,7 +191,7 @@ def sample(
                 generator=generator,
             )
         case ("temperature", temperature):
-            tokens, softmax = temperature_sampling_batch(
+            tokens, softmax = top_k_top_p_sampling_batch(
                 logits,
                 temperature=cast(float, temperature),
                 generator=generator,
@@ -296,8 +286,8 @@ def _sample_from_probs(
             probs: torch.Tensor,
             generator: Optional[torch.Generator],
         ) -> torch.Tensor:
-            return sampling_from_probs_generator_op(
-                probs, generator, check_nan=cls._flashinfer_check_nans(probs)
+            return sampling_from_probs_op(
+                probs, generator=generator, check_nan=cls._flashinfer_check_nans(probs)
             )
 
         def _sample_greedy_with_probs(
@@ -556,11 +546,11 @@ def sample(
             logits = self._prepare_logits_with_temperature(
                 logits, group_logit_indices, self._temperature
             )
-            return top_k_top_p_sampling_from_logits_with_generator_op(
+            return top_k_top_p_sampling_from_logits_op(
                 logits,
                 self._top_k,
                 self._top_p,
-                generator,
+                generator=generator,
                 check_nan=self._flashinfer_check_nans(logits),
             ), None
 
@@ -591,8 +581,11 @@ def sample(
             probs = self._prepare_probs_with_temperature(
                 logits, group_logit_indices, self._temperature
             )
-            return top_k_sampling_from_probs_generator_op(
-                probs, self._top_k, generator, check_nan=self._flashinfer_check_nans(probs)
+            return top_k_sampling_from_probs_op(
+                probs,
+                self._top_k,
+                generator=generator,
+                check_nan=self._flashinfer_check_nans(probs),
             ), None
 
     class TopPSampleOnly(StrategyImplSampleOnly):
@@ -622,8 +615,11 @@ def sample(
             probs = self._prepare_probs_with_temperature(
                 logits, group_logit_indices, self._temperature
             )
-            return top_p_sampling_from_probs_generator_op(
-                probs, self._top_p, generator, check_nan=self._flashinfer_check_nans(probs)
+            return top_p_sampling_from_probs_op(
+                probs,
+                self._top_p,
+                generator=generator,
+                check_nan=self._flashinfer_check_nans(probs),
             ), None
 
     class TemperatureOnlySampleOnly(StrategyImplSampleOnly):
@@ -814,7 +810,7 @@ def sample_grouped_strategies(
 
 
 # Re-export the torch greedy op (used by drafting_loops and speculative/interface).
-greedy = _torch_greedy
+greedy = greedy_search_sampling_batch
 
 # ---------------------------------------------------------------------------
 # Spec-decoding interface: compute_probs_from_logits (per-request tensor params)
@@ -825,12 +821,11 @@ def sanitize_top_k(top_k: torch.Tensor, vocab_size: int) -> torch.Tensor:
     """Map ``top_k`` into a backend-safe range before top-k filtering.
 
     Per ``SamplingParams``, ``top_k == 0`` means "all logits" (top-k disabled),
-    but the flashinfer (``top_k_mask_logits``) and PyTorch-native top-k paths
-    break on a literal 0 — they mask the entire row (all-zero probs) or gather
-    out of bounds. Mirror the C++ op (``dynamicTreeKernels.cu``): map any
-    non-positive value (and any oversized disable sentinel such as
-    ``INT32_MAX``) to ``vocab_size`` (== keep all tokens), leaving genuine
-    top_k values untouched.
+    but the flashinfer top-k kernels (``top_k_mask_logits``) break on a literal
+    0 — they mask the entire row (all-zero probs). Map any non-positive value
+    (and any oversized disable sentinel such as ``INT32_MAX``) to
+    ``vocab_size`` (== keep all tokens), leaving genuine top_k values
+    untouched.
     """
     return torch.where(top_k > 0, top_k, torch.full_like(top_k, vocab_size)).clamp(max=vocab_size)
 
@@ -859,8 +854,8 @@ def sampling_batch_spec_dec_one_model(
     temperatures: torch.Tensor,
     top_k: torch.Tensor,
     top_p: torch.Tensor,
-    seed: Optional[int] = None,
-    offset: Optional[int] = None,
+    seed: Optional[torch.Tensor] = None,
+    offset: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """CUDA-graph compatible sampling; supports mixed sampling params. Returns sampled tokens."""
     top_k = sanitize_top_k(top_k, logits.shape[-1])

From c6f708b8b93d3c7f81592196461ccabdb1657dba Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <zhaoyangw@nvidia.com>
Date: Mon, 13 Jul 2026 03:15:08 -0700
Subject: [PATCH 2/5] [TRTLLM-13212][refactor] Promote cross-module sampler op
 internals to public API

The sampling facade (sampling_utils) and sampler.py were reaching into
underscore-private members of ops.vanilla. Promote the de-facto public
surface and remove a duplicated constant:

- _GREEDY_TEMPERATURE_THRESHOLD -> GREEDY_TEMPERATURE_THRESHOLD; document
  the contract with the spec-decoding metadata layer, and derive
  DISABLE_TEMP_VAL in speculative/interface.py from it
  (GREEDY_TEMPERATURE_THRESHOLD / 10 == 1e-5, value unchanged) so the
  sentinel and the threshold can no longer drift apart silently.
- _safely_apply_temperature_inplace -> safely_apply_temperature_inplace.
- _Fusions -> Fusions (compiled logprobs fusion helpers used by
  sampler.py); per-method _impl internals stay private.
- Drop the duplicated BEAM_SEARCH_PAD_TOKEN definition in sampling_utils;
  ops.vanilla now holds the single definition and the facade re-exports it.

speculative/interface.py imports the threshold via the sampling_utils
facade, keeping sampling_utils the only importer of the ops backend
modules.

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
---
 .../_torch/pyexecutor/sampler/ops/vanilla.py  | 24 +++++++++++--------
 .../_torch/pyexecutor/sampler/sampler.py      |  8 +++----
 .../pyexecutor/sampler/sampling_utils.py      | 14 +++++++----
 tensorrt_llm/_torch/speculative/interface.py  | 10 ++++++--
 4 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
index 1232ed37ea4e..5307876ae61d 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
@@ -26,7 +26,7 @@
 from tensorrt_llm._utils import prefer_pinned
 from tensorrt_llm.bindings.executor import FinishReason
 
-_BEAM_SEARCH_PAD_TOKEN = -1
+BEAM_SEARCH_PAD_TOKEN = -1
 
 
 @dataclass(kw_only=True)
@@ -235,7 +235,7 @@ def beam_search_sampling_batch(
 
     next_tokens = next_tokens % vocab_size
     ended_predecessor_mask = torch.gather(dim=1, index=predecessor_beam, input=finished_beams_mask)
-    next_tokens = torch.where(ended_predecessor_mask, _BEAM_SEARCH_PAD_TOKEN, next_tokens)
+    next_tokens = torch.where(ended_predecessor_mask, BEAM_SEARCH_PAD_TOKEN, next_tokens)
 
     old_cum_log_probs = beam_search_args.cum_log_probs[beam_search_args.seq_slots].view(-1)
     beam_search_args.new_log_probs[beam_search_args.seq_slots, :beam_width_out] = (
@@ -283,15 +283,19 @@ def sample_rejected(
     return cast(int, new_token.item())
 
 
-_GREEDY_TEMPERATURE_THRESHOLD = 1e-4
+# Rows whose temperature is at or below this threshold are treated as greedy.
+# Contract with the spec-decoding metadata layer: greedy requests are
+# normalized to a sentinel temperature strictly below this threshold (see
+# DISABLE_TEMP_VAL in speculative/interface.py, which derives from it).
+GREEDY_TEMPERATURE_THRESHOLD = 1e-4
 
 
-def _safely_apply_temperature_inplace(
+def safely_apply_temperature_inplace(
     logits_inout: torch.Tensor, temp: torch.Tensor
 ) -> torch.Tensor:
     """Divide logits by per-row temperature in place, guarding the greedy sentinel.
 
-    Greedy requests carry a temperature of 0 / <= ``_GREEDY_TEMPERATURE_THRESHOLD``.
+    Greedy requests carry a temperature of 0 / <= ``GREEDY_TEMPERATURE_THRESHOLD``.
     Dividing by it would blow logits up to inf/nan and corrupt downstream sampling
     (argmax / softmax / multinomial). Those rows are clamped to a temperature of 1.0
     so the division is numerically safe; callers are expected to overwrite the greedy
@@ -301,11 +305,11 @@ def _safely_apply_temperature_inplace(
     ``logits_inout`` is modified in place (``div_``) and also returned for
     convenience; ``temp`` is left untouched.
     """
-    safe_temp = torch.where(temp <= _GREEDY_TEMPERATURE_THRESHOLD, torch.ones_like(temp), temp)
+    safe_temp = torch.where(temp <= GREEDY_TEMPERATURE_THRESHOLD, torch.ones_like(temp), temp)
     return logits_inout.div_(safe_temp.unsqueeze(dim=1))
 
 
-class _Fusions:
+class Fusions:
     @staticmethod
     @torch.compile(dynamic=None, fullgraph=True)
     def _gather_scatter_impl(
@@ -327,7 +331,7 @@ def gather_scatter(
         torch._dynamo.mark_dynamic(dst_index_cuda, 0)
         torch._dynamo.mark_dynamic(src_cuda, 0)
         torch._dynamo.mark_dynamic(src_index_cuda, 0)
-        _Fusions._gather_scatter_impl(dst_cuda, dst_index_cuda, src_cuda, src_index_cuda)
+        Fusions._gather_scatter_impl(dst_cuda, dst_index_cuda, src_cuda, src_index_cuda)
 
     @staticmethod
     @torch.compile(dynamic=None, fullgraph=True)
@@ -345,7 +349,7 @@ def determine_sampled_rank(
     ) -> torch.Tensor:
         torch._dynamo.mark_dynamic(group_logprobs_cuda, 0)
         torch._dynamo.mark_dynamic(sampled_logprobs_cuda, 0)
-        return _Fusions._determine_sampled_rank_impl(group_logprobs_cuda, sampled_logprobs_cuda)
+        return Fusions._determine_sampled_rank_impl(group_logprobs_cuda, sampled_logprobs_cuda)
 
     @staticmethod
     @torch.compile(
@@ -368,4 +372,4 @@ def _gather_log_softmax_impl(
     def gather_log_softmax(inputs_cuda: torch.Tensor, indices_cuda: torch.Tensor) -> 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)
+        return Fusions._gather_log_softmax_impl(inputs_cuda, indices_cuda)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
index e7fe112ace6f..2a2a1f3edb2a 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
@@ -93,11 +93,11 @@
     GREEDY,
     BeamSearchMetadata,
     FlashInferGroupedStrategySampler,
+    Fusions,
     GenericStrategyKeyType,
     Strategy,
     StrategyMetadata,
     UtilsSamplingParams,
-    _Fusions,
     get_rejected_indices,
     resolve_sampling_strategy,
     sample,
@@ -4217,7 +4217,7 @@ def _sample_batched_by_strategy(
                 logit_indices_for_raw_logprobs_cuda += batch_next_tokens_offset_start
                 # NB: Copy could be avoided by storing logit indices (and temperature) instead (cf. comment on
                 #     processed logprobs above).
-                _Fusions.gather_scatter(
+                Fusions.gather_scatter(
                     batch_logits_for_logprobs_cuda,
                     logit_indices_for_raw_logprobs_cuda,
                     group_logits_cuda,
@@ -4563,7 +4563,7 @@ def _process_logprobs(
             )
 
             # (batch_size, vocab_size)
-            group_logprobs_cuda = _Fusions.gather_log_softmax(
+            group_logprobs_cuda = Fusions.gather_log_softmax(
                 batched_sampling_result.batch_logits_for_logprobs_cuda, group_logits_indices_cuda
             )
 
@@ -4603,7 +4603,7 @@ def _process_logprobs(
             # NB: Computation of sampled rank could be lowered into FlashInferGroupedStrategySampler, s.t., e.g., for
             #     greedy sampling, logits management and log_softmax could be completely skipped (sampled rank
             #     computation is trivial in this case).
-            sampled_rank_cuda = _Fusions.determine_sampled_rank(
+            sampled_rank_cuda = Fusions.determine_sampled_rank(
                 group_logprobs_cuda, sampled_vals_cuda
             )
 
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
index f27e9bb33b7b..da8f6e220c31 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
@@ -49,11 +49,14 @@
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
     top_p_sampling_from_probs_op as top_p_sampling_from_probs_op,
 )
+from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
+    GREEDY_TEMPERATURE_THRESHOLD as GREEDY_TEMPERATURE_THRESHOLD,
+)
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
     BeamSearchMetadata as BeamSearchMetadata,
 )
+from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import Fusions as Fusions
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import StrategyMetadata as StrategyMetadata
-from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import _Fusions as _Fusions
 from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import (
     beam_search_sampling_batch as beam_search_sampling_batch,
 )
@@ -86,7 +89,8 @@
 
 Strategy: TypeAlias = TopK | TopP | Greedy | TopKTopP | TemperatureOnly | BeamSearch
 
-BEAM_SEARCH_PAD_TOKEN = -1
+# Re-exported from the beam-search op implementation (single source of truth).
+BEAM_SEARCH_PAD_TOKEN = vanilla.BEAM_SEARCH_PAD_TOKEN
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -861,13 +865,13 @@ def sampling_batch_spec_dec_one_model(
     top_k = sanitize_top_k(top_k, logits.shape[-1])
     # Greedy rows (temperature <= threshold) must return the argmax token, not a
     # sample from the temperature-scaled distribution. Capture the argmax from the
-    # *original* logits up front; _safely_apply_temperature_inplace then guards the division
+    # *original* logits up front; safely_apply_temperature_inplace then guards the division
     # against the greedy sentinel, and torch.where restores the greedy rows below.
     # All ops are branch-free (no data-dependent control flow), so this stays
     # CUDA-graph safe.
-    is_greedy = temperatures <= vanilla._GREEDY_TEMPERATURE_THRESHOLD
+    is_greedy = temperatures <= vanilla.GREEDY_TEMPERATURE_THRESHOLD
     greedy_tokens = logits.argmax(dim=-1)
-    logits = vanilla._safely_apply_temperature_inplace(logits, temperatures)
+    logits = vanilla.safely_apply_temperature_inplace(logits, temperatures)
     sampled = flashinfer.top_k_top_p_sampling_from_logits_op(
         logits, top_k, top_p, seed=seed, offset=offset
     )
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index d0b0db6e3d32..41cab607e396 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -634,10 +634,16 @@ def _scan_one_model_sampling(
         before the CUDA graph key is built.
         """
         from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
+        from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import \
+            GREEDY_TEMPERATURE_THRESHOLD
         from tensorrt_llm.sampling_params import SamplingParams
 
-        # Need to use a very small value for temperature when disabled to avoid division by 0
-        DISABLE_TEMP_VAL = 1e-5
+        # Sentinel temperature for greedy / temperature-disabled rows. Must stay
+        # strictly below GREEDY_TEMPERATURE_THRESHOLD so the sampling kernels
+        # recognize these rows as greedy; small enough that even if a row were
+        # (incorrectly) sampled, softmax(logits / val) is effectively one-hot,
+        # and non-zero to avoid division by 0.
+        DISABLE_TEMP_VAL = GREEDY_TEMPERATURE_THRESHOLD / 10
         # Very large values disable topk.
         DISABLE_TOPK_VAL = torch.iinfo(torch.int32).max
         DISABLE_TOPP_VAL = 1.0

From e45fe40bbd284e9137f530b80a62819338adcaa1 Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <zhaoyangw@nvidia.com>
Date: Mon, 13 Jul 2026 04:07:53 -0700
Subject: [PATCH 3/5] [TRTLLM-13212][chore] Remove dead TorchSampler methods

Repo-wide audit found four methods with zero callers (no dynamic
dispatch, no test usage):

- _handle_finish_reasons: superseded by _handle_first_finish_reasons /
  direct _handle_finish_reasons_impl calls.
- _longest_stop_word_len: logic re-implemented in the live
  _check_stop_words_length.
- _requests_with_stop_words / _request_indices_with_stop_words: unused.

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
---
 .../_torch/pyexecutor/sampler/sampler.py      | 56 -------------------
 1 file changed, 56 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
index 2a2a1f3edb2a..5ed854cf0706 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
@@ -2544,31 +2544,6 @@ def _handle_finish_reasons_impl(
             return True
         return False
 
-    def _handle_finish_reasons(
-        self,
-        request: LlmRequest,
-        finish_reasons: torch.Tensor,
-        finish_reasons_list: list[list[list[int]]],
-    ) -> bool:
-        """Check if all beams of a request have finished and set the request state accordingly
-
-        Args:
-            request: LlmRequest. The request to check.
-            finish_reasons: torch.Tensor. Shape: (max_tokens, max_batch_size, max_beam_width)
-                            The finish reasons for each beam.
-            finish_reasons_list: list[list[list[int]]]. The finish reasons for each beam.
-        Returns:
-            True if all beams have finished, False otherwise.
-        """
-        assert request.py_seq_slot is not None
-        beam_width = request.py_beam_width
-        return self._handle_finish_reasons_impl(
-            request,
-            beam_width,
-            finish_reasons[DEFAULT_STEP_IDX, request.py_seq_slot],
-            finish_reasons_list[request.py_seq_slot][DEFAULT_STEP_IDX],
-        )
-
     def _handle_first_finish_reasons(
         self,
         request: LlmRequest,
@@ -4466,37 +4441,6 @@ def _select_generated_logits(
 
         return sampling_requests, sampling_requests_metadata, logits_cuda
 
-    @staticmethod
-    def _longest_stop_word_len(requests: Iterable[LlmRequest]) -> int:
-        max_stop_word_len = 0
-        for req in requests:
-            assert req.py_stop_words_list is not None
-            _, cumsum = req.py_stop_words_list
-            if -1 in cumsum:
-                cumsum = cumsum[: cumsum.index(-1)]
-            request_max_stop_word_len = np.max(np.diff(cumsum, prepend=0), initial=0).item()
-            max_stop_word_len = max(max_stop_word_len, request_max_stop_word_len)
-        return max_stop_word_len
-
-    @staticmethod
-    def _requests_with_stop_words(requests: list[LlmRequest]) -> list[LlmRequest]:
-        return [
-            r
-            for r in requests
-            if (r.py_stop_words_list is not None and len(r.py_stop_words_list[0]) > 0)
-        ]
-
-    def _request_indices_with_stop_words(self, requests: list[LlmRequest]) -> torch.Tensor:
-        return torch.tensor(
-            [
-                ridx
-                for ridx, r in enumerate(requests)
-                if (r.py_stop_words_list is not None and len(r.py_stop_words_list[0]) > 0)
-            ],
-            dtype=torch.int32,
-            pin_memory=prefer_pinned(),
-        ).to(device="cuda", non_blocking=True)
-
     @nvtx_range("_process_logprobs")
     def _process_logprobs(
         self,

From 5cf0c6adabb2d28ba4be3c1afc5f22f3208b28b3 Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <zhaoyangw@nvidia.com>
Date: Tue, 14 Jul 2026 01:02:28 -0700
Subject: [PATCH 4/5] [TRTLLM-13212][chore] Fix logprobs history slicing,
 simplify min-length penalty path, document flashinfer op contracts

- _get_logprobs_from_request: use forward slicing for the history view;
  [:, :-preallocate_extra_steps, :] silently yields an empty view when
  preallocate_extra_steps == 0 (the default), skipping history filling.
- _apply_min_length_penalty: return None instead of echoing the
  in-place-mutated logits; defer the num_steps/num_beams host
  tensor->list conversion until after the early-out; flatten the
  double-nested guard ifs into early continues, hoisting the per-request
  offset accumulation above the guards.
- _flashinfer_check_nans: restore the explanatory comments lost in the
  ops-module split (#15542): the constant False deliberately keeps
  FlashInfer's syncing check_nan path disabled, while the async
  device-side assert provides NaN protection without stalling the
  pipeline (see flashinfer issue #1575).
- ops/flashinfer: per-op docstrings pointing at the module-level
  generator-vs-seed/offset contract; note why the 1:1 pipeline-stage
  wrappers (mask/softmax/renorm) exist (guarded import, PDL decision).

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
---
 .../pyexecutor/sampler/ops/flashinfer.py      | 26 +++++++
 .../_torch/pyexecutor/sampler/sampler.py      | 69 +++++++++++--------
 .../pyexecutor/sampler/sampling_utils.py      | 11 +++
 3 files changed, 76 insertions(+), 30 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
index 4d8226bc4673..0dc72d6cb158 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
@@ -49,6 +49,11 @@ def top_k_top_p_sampling_from_logits_op(
     offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
+    """Fused top-k + top-p sampling from pre-softmax logits.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
         logits,
         top_k=top_k,
@@ -71,6 +76,11 @@ def sampling_from_probs_op(
     offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
+    """Categorical sampling from probabilities.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
         probs,
         deterministic=True,
@@ -91,6 +101,11 @@ def top_k_sampling_from_probs_op(
     offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
+    """Top-k filtered sampling from probabilities.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.top_k_sampling_from_probs(
         probs,
         top_k=top_k,
@@ -112,6 +127,11 @@ def top_p_sampling_from_probs_op(
     offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
+    """Top-p filtered sampling from probabilities.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.top_p_sampling_from_probs(
         probs,
         top_p=top_p,
@@ -124,6 +144,12 @@ def top_p_sampling_from_probs_op(
     return tokens
 
 
+# The three ops below wrap the mask -> softmax -> renorm pipeline stages 1:1.
+# The wrappers exist so callers stay importable without flashinfer installed
+# (the flashinfer import above is guarded); softmax_op additionally centralizes
+# the PDL env decision.
+
+
 def softmax_op(
     logits: torch.Tensor,
     temperature: Optional[torch.Tensor],
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
index 5ed854cf0706..27323d0b846a 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
@@ -3174,8 +3174,10 @@ def _get_logprobs_from_request(
             pin_memory=pin_memory,
             dtype=torch.int32,
         )
-        logprobs_tensor = logprobs_tensor_full[:, :-preallocate_extra_steps, :]
-        logprobs_indices_tensor = logprobs_indices_tensor_full[:, :-preallocate_extra_steps, :]
+        # NB: forward slicing, because [:, :-0, :] would yield an empty view
+        #     instead of the full history when preallocate_extra_steps == 0.
+        logprobs_tensor = logprobs_tensor_full[:, :num_generated_tokens, :]
+        logprobs_indices_tensor = logprobs_indices_tensor_full[:, :num_generated_tokens, :]
         if logprobs_tensor.numel() > 0:
             logprobs_list = request.py_result.log_probs
             assert logprobs_list is not None
@@ -4278,45 +4280,54 @@ def _dims_canonically_ordered(t: torch.Tensor) -> bool:
     def _apply_min_length_penalty(
         logits: torch.Tensor,
         requests: list[LlmRequest],
-        num_steps: list[int],
-        num_beams: list[int],
-    ) -> torch.Tensor:
-        """Inplace apply min_length_penalty to logits.
+        num_steps_tensor: torch.Tensor,
+        num_beams_tensor: torch.Tensor,
+    ) -> None:
+        """Apply min_length_penalty to logits, mutating ``logits`` in place.
 
         Args:
             logits: The logits to apply min length penalty to
             requests: The requests to apply min length penalty to
-            num_steps: The number of steps per request
-
-        Returns:
-            The logits with min length penalty applied
+            num_steps_tensor: The number of steps per request (host tensor)
+            num_beams_tensor: The number of beams per request (host tensor)
         """
         if not any(
             r.py_min_length and (r.max_beam_num_tokens - r.py_orig_prompt_len) < r.py_min_length[0]
             for r in requests
         ):
-            return logits
+            return
+
+        # Deferred host conversion: only needed on the (rare) penalty path.
+        num_steps = num_steps_tensor.tolist()
+        num_beams = num_beams_tensor.tolist()
 
         rows: list[int] = []
         cols: list[int] = []
         current_offset = 0
         for index, r in enumerate(requests):
-            if r.py_min_length:
-                # Use the original end_id (before ignore_eos override)
-                # so we suppress the real EOS token, not token -1.
-                end_id = getattr(r, "py_original_end_id", r.py_end_id)
-                if end_id is not None and end_id > -1:
-                    for beam_idx in range(num_beams[index]):
-                        for step in range(num_steps[index]):
-                            if (
-                                r.get_num_tokens(beam_idx) - r.py_orig_prompt_len
-                            ) + step < r.py_min_length[0]:
-                                rows.append(current_offset + num_steps[index] * beam_idx + step)
-                                cols.append(end_id)
-                            else:
-                                break
+            # Advance the offset before any guard below can skip the request:
+            # every request occupies its logits rows, penalized or not.
+            req_offset = current_offset
             current_offset += num_steps[index] * num_beams[index]
 
+            if not r.py_min_length:
+                continue
+            # Use the original end_id (before ignore_eos override)
+            # so we suppress the real EOS token, not token -1.
+            end_id = getattr(r, "py_original_end_id", r.py_end_id)
+            if end_id is None or end_id <= -1:
+                continue
+
+            for beam_idx in range(num_beams[index]):
+                for step in range(num_steps[index]):
+                    if (r.get_num_tokens(beam_idx) - r.py_orig_prompt_len) + step < r.py_min_length[
+                        0
+                    ]:
+                        rows.append(req_offset + num_steps[index] * beam_idx + step)
+                        cols.append(end_id)
+                    else:
+                        break
+
         if rows:
             neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device)
             row_idx = torch.tensor(rows, dtype=torch.long, pin_memory=prefer_pinned()).to(
@@ -4327,8 +4338,6 @@ def _apply_min_length_penalty(
             )
             logits.index_put_((row_idx, col_idx), neg_inf, accumulate=False)
 
-        return logits
-
     @staticmethod
     def _select_generated_logits(
         scheduled_requests: ScheduledRequests,
@@ -4634,11 +4643,11 @@ def _process_requests(
             logits_cuda, sampling_requests, sampling_requests_metadata.req_num_steps
         )
 
-        logits_cuda = self._apply_min_length_penalty(
+        self._apply_min_length_penalty(
             logits_cuda,
             sampling_requests,
-            sampling_requests_metadata.req_num_steps.tolist(),
-            sampling_requests_metadata.req_num_beams.tolist(),
+            sampling_requests_metadata.req_num_steps,
+            sampling_requests_metadata.req_num_beams,
         )
 
         # Fast path for greedy sampling
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
index da8f6e220c31..fc93d622b64a 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
@@ -249,8 +249,19 @@ def sample(
         ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
             pass
 
+        # TODO: Revisit this after determining performance impact
+        #
+        # NB: NaN logits can lead to crashes, see
+        #     https://github.com/flashinfer-ai/flashinfer/issues/1575
+        #
         @staticmethod
         def _flashinfer_check_nans(inputs: torch.Tensor) -> bool:
+            # Deliberately returns False to keep FlashInfer's own 'check_nan' path
+            # disabled: that path is a host-side `if torch.any(torch.isnan(...))`,
+            # which forces a device sync on every call. The explicit async
+            # device-side assert below provides the same protection without
+            # stalling the pipeline.
+            # https://github.com/pytorch/pytorch/issues/36853
             torch._assert_async(~torch.any(torch.isnan(inputs)))
             return False
 

From 003a19f62571f6c22a16b201c53845e04b9b24db Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <zhaoyangw@nvidia.com>
Date: Tue, 14 Jul 2026 19:42:01 -0700
Subject: [PATCH 5/5] [TRTLLM-13212][fix] Accept seed/offset in flashinfer
 sampling test mocks

The unified flashinfer ops now forward seed/offset (alongside generator) to
the underlying flashinfer.sampling.* calls, but the test mocks that patch
those functions only accepted generator, so the eager batched-sampling tests
failed with 'unexpected keyword argument seed'. Add seed/offset (defaulting to
None) to the four affected mock signatures.

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
---
 tests/unittest/_torch/sampler/test_torch_sampler.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py
index 5bda91081396..c16da7a2c70e 100644
--- a/tests/unittest/_torch/sampler/test_torch_sampler.py
+++ b/tests/unittest/_torch/sampler/test_torch_sampler.py
@@ -1894,6 +1894,8 @@ def _mock_flashinfer_top_k_top_p(
             deterministic: bool,
             check_nan: bool,
             generator: torch.Generator,
+            seed: Optional[Union[int, torch.Tensor]] = None,
+            offset: Optional[Union[int, torch.Tensor]] = None,
         ) -> torch.Tensor:
             assert filter_apply_order == "top_k_first"
             assert deterministic
@@ -1960,6 +1962,8 @@ def _mock_flashinfer_top_k(
             deterministic: bool,
             check_nan: bool,
             generator: torch.Generator,
+            seed: Optional[Union[int, torch.Tensor]] = None,
+            offset: Optional[Union[int, torch.Tensor]] = None,
         ) -> torch.Tensor:
             assert deterministic
             assert not check_nan, "check_nan syncs"
@@ -1991,6 +1995,8 @@ def _mock_flashinfer_top_p(
             deterministic: bool,
             check_nan: bool,
             generator: torch.Generator,
+            seed: Optional[Union[int, torch.Tensor]] = None,
+            offset: Optional[Union[int, torch.Tensor]] = None,
         ) -> torch.Tensor:
             assert deterministic
             assert not check_nan, "check_nan syncs"
@@ -2021,6 +2027,8 @@ def _mock_flashinfer_from_probs(
             deterministic: bool,
             check_nan: bool,
             generator: torch.Generator,
+            seed: Optional[Union[int, torch.Tensor]] = None,
+            offset: Optional[Union[int, torch.Tensor]] = None,
         ) -> torch.Tensor:
             assert deterministic
             assert not check_nan, "check_nan syncs"