fix https://github.com/NVIDIA/TensorRT-Edge-LLM/issues/211: advance the Philox offset per sampled token

cpp/runtime/decoding/vanillaDecoder.cpp

@@ -37,6 +37,12 @@ namespace trt_edgellm
 {
 namespace rt
 {
+namespace
+{
+//! Fixed seed keeps a given input reproducible; the offset supplies the per-step variation.
+constexpr uint64_t kSAMPLING_PHILOX_SEED{42};
+} // namespace
+
 namespace
 {
 constexpr int32_t kDecodeProfile{1};
@@ -121,8 +127,11 @@ bool VanillaDecoder::decodeStep(DecodingInferenceContext& context)
     {
         SamplingParams params(activeBatchSize, mRuntime.deployment.base.outputVocabSize, context.temperature,
             static_cast<int32_t>(context.topK), context.topP);
+        // Advance the Philox offset per sampled token. With the default offset of 0 every
+        // step shares one RNG counter, so the same uniform is drawn each time and sampling
+        // degenerates to greedy regardless of temperature and top_p.
         topKtopPSamplingFromLogits(mRuntime.base.pipelineIO.outputLogits, mRuntime.sampling.indices, params,
-            mRuntime.sampling.workspace, context.stream);
+            mRuntime.sampling.workspace, context.stream, kSAMPLING_PHILOX_SEED, mSamplingPhiloxOffset++);
     }
     else
     {

cpp/runtime/decoding/vanillaDecoder.h

@@ -67,11 +67,23 @@ class VanillaDecoder final : public DecodingStrategy
     {
     }
 
-    void resetForNewSequences(Tensor&, cudaStream_t) override {}
+    void resetForNewSequences(Tensor&, cudaStream_t) override
+    {
+        //! Restart the sampler's RNG stream so a given input reproduces exactly.
+        mSamplingPhiloxOffset = 0;
+    }
     void onBatchEvict(std::vector<int32_t> const&, int32_t, int32_t, Tensor&, cudaStream_t) override {}
 
 private:
     DecodingRuntimeContext& mRuntime;
+
+    //! Philox offset for top-k/top-p sampling, advanced once per sampled token.
+    //!
+    //! curand_init(seed, batchIdx, offset) is keyed on the offset, so leaving it at the
+    //! default 0 draws the SAME uniform at every decode step: sampling becomes
+    //! deterministic and collapses onto the argmax, making temperature and top_p inert.
+    //! The TTS talker path already varies this deliberately (qwen3OmniTTSRuntime.cpp).
+    uint64_t mSamplingPhiloxOffset{0};
 };
 
 } // namespace rt

cpp/runtime/llmRankRuntime.cpp

@@ -71,6 +71,11 @@ constexpr int32_t kDecodeProfile{1};
 
 namespace rt
 {
+namespace
+{
+//! Fixed seed keeps a given input reproducible; the per-call offset supplies the variation.
+constexpr uint64_t kSAMPLING_PHILOX_SEED{42};
+} // namespace
 
 std::vector<int32_t> LLMRankRuntime::countPromptTokens(LLMGenerationRequest const& request) const
 {
@@ -2578,8 +2583,8 @@ bool LLMRankRuntime::runBaseModelPrefill(
     {
         SamplingParams params(activeBatchSize, mDeployment.base.outputVocabSize, context.temperature,
             static_cast<int32_t>(context.topK), context.topP);
-        topKtopPSamplingFromLogits(
-            mPipelineIO->outputLogits, mSamplingIndices, params, mSamplingWorkspace, context.stream);
+        topKtopPSamplingFromLogits(mPipelineIO->outputLogits, mSamplingIndices, params, mSamplingWorkspace,
+            context.stream, kSAMPLING_PHILOX_SEED, mSamplingPhiloxOffset++);
     }
     else
     {

cpp/runtime/llmRankRuntime.h

@@ -320,6 +320,11 @@ class LLMRankRuntime
     // [2] Sampling workspace and output tensors that used across all the sampling operations.
     rt::Tensor mSamplingWorkspace;
     rt::Tensor mSamplingIndices;
+
+    //! Philox offset for top-k/top-p sampling, advanced once per sampled token.
+    //! See the note on VanillaDecoder::mSamplingPhiloxOffset: a fixed offset makes every
+    //! sampling call draw the same uniform, which collapses sampling onto the argmax.
+    uint64_t mSamplingPhiloxOffset{0};
     rt::Tensor mSamplingScores;
     rt::Tensor mBaseVocabMappingTable; // Vocab mapping table for base model reduced vocab (empty if not used)
 

unittests/cpp/sampler/samplingTests.cpp

@@ -308,6 +308,53 @@ TEST_F(SamplingTest, TemperatureZeroParameterOverride)
     }
 }
 
+TEST_F(SamplingTest, AdvancingPhiloxOffsetChangesSampledToken)
+{
+    // Pins the Philox offset contract that #211 depended on: distinct offsets must be able
+    // to select distinct tokens, and a repeated offset must reproduce.
+    //
+    // SCOPE, stated so this is not mistaken for a regression guard: the kernel was never the
+    // faulty part. #211 was two CALLERS that never passed an offset, so it defaulted to 0 at
+    // every decode step and one uniform was drawn for the whole sequence. This test passes
+    // on the unpatched tree. It documents the invariant a caller has to uphold; catching a
+    // caller that does not would need a decoder-level test with a real runtime.
+    constexpr int32_t kBatchSize = 1;
+    constexpr int32_t kVocabSize = 32;
+    constexpr int32_t kNumOffsets = 16; // each call carries a device sync; 16 is enough to diverge
+
+    // A deliberately flat distribution: with many near-equal candidates, a working sampler
+    // visits several of them while a broken one returns the argmax every time.
+    std::vector<float> hostLogits(kVocabSize, 1.0f);
+
+    rt::Tensor logits({kBatchSize, kVocabSize}, rt::DeviceType::kGPU, nvinfer1::DataType::kFLOAT);
+    CUDA_CHECK(
+        cudaMemcpy(logits.rawPointer(), hostLogits.data(), hostLogits.size() * sizeof(float), cudaMemcpyHostToDevice));
+
+    rt::Tensor selected({kBatchSize, 1}, rt::DeviceType::kGPU, nvinfer1::DataType::kINT32);
+    SamplingParams params(kBatchSize, kVocabSize, /*temperature=*/1.0f, /*topK=*/kVocabSize, /*topP=*/1.0f);
+    size_t const workspaceSize = getTopKtopPSamplingWorkspaceSize(kBatchSize, kVocabSize, params);
+    rt::Tensor workspace({static_cast<int64_t>(workspaceSize)}, rt::DeviceType::kGPU, nvinfer1::DataType::kINT8);
+
+    auto sampleAtOffset = [&](uint64_t offset) {
+        topKtopPSamplingFromLogits(logits, selected, params, workspace, 0, TEST_SEED, offset);
+        CUDA_CHECK(cudaDeviceSynchronize());
+        return copyDeviceToHost<int32_t>(selected).at(0);
+    };
+
+    std::set<int32_t> distinctTokens;
+    for (uint64_t offset = 0; offset < kNumOffsets; ++offset)
+    {
+        distinctTokens.insert(sampleAtOffset(offset));
+    }
+
+    // The bug produced exactly one distinct token across every offset.
+    EXPECT_GT(distinctTokens.size(), 1U) << "sampling returned the same token for all " << kNumOffsets
+                                         << " Philox offsets, so the offset is not reaching the RNG";
+
+    // Same seed and same offset must still be reproducible.
+    EXPECT_EQ(sampleAtOffset(7), sampleAtOffset(7));
+}
+
 TEST(SamplingUtilsTest, ShouldUseNonGreedySampling)
 {
     EXPECT_FALSE(trt_edgellm::shouldUseNonGreedySampling(1.0f, 0, 1.0f));