diff --git a/cpp/runtime/decoding/vanillaDecoder.cpp b/cpp/runtime/decoding/vanillaDecoder.cpp index a69de91a..f7d88e2e 100644 --- a/cpp/runtime/decoding/vanillaDecoder.cpp +++ b/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(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 { diff --git a/cpp/runtime/decoding/vanillaDecoder.h b/cpp/runtime/decoding/vanillaDecoder.h index 454acdbf..ad5f8256 100644 --- a/cpp/runtime/decoding/vanillaDecoder.h +++ b/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 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 diff --git a/cpp/runtime/llmRankRuntime.cpp b/cpp/runtime/llmRankRuntime.cpp index 9548d564..6aaa3fa7 100644 --- a/cpp/runtime/llmRankRuntime.cpp +++ b/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 LLMRankRuntime::countPromptTokens(LLMGenerationRequest const& request) const { @@ -2578,8 +2583,8 @@ bool LLMRankRuntime::runBaseModelPrefill( { SamplingParams params(activeBatchSize, mDeployment.base.outputVocabSize, context.temperature, static_cast(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 { diff --git a/cpp/runtime/llmRankRuntime.h b/cpp/runtime/llmRankRuntime.h index 984caff5..2628b084 100644 --- a/cpp/runtime/llmRankRuntime.h +++ b/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) diff --git a/unittests/cpp/sampler/samplingTests.cpp b/unittests/cpp/sampler/samplingTests.cpp index 08ebb18c..f2720593 100644 --- a/unittests/cpp/sampler/samplingTests.cpp +++ b/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 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(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(selected).at(0); + }; + + std::set 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));