From 4a3892cd279ef2129182d9214fdb3a72b71e4dee Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Sun, 29 Jun 2025 13:56:37 +0000 Subject: [PATCH 1/4] refactor: Simplify disableLookahead method in DecoderState - Removed the RequestVector parameter from the disableLookahead method in DecoderState, streamlining its interface. - Updated calls to disableLookahead in TrtGptModelInflightBatching to reflect the new method signature. - Adjusted internal logic to maintain functionality without the need for external request data. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/decoderState.h | 2 +- .../batch_manager/trtGptModelInflightBatching.cpp | 8 +++++++- cpp/tensorrt_llm/runtime/decoderState.cpp | 11 +---------- 3 files changed, 9 insertions(+), 12 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/decoderState.h b/cpp/include/tensorrt_llm/runtime/decoderState.h index 95d7ff0ffac9..ca1c8a7cc740 100644 --- a/cpp/include/tensorrt_llm/runtime/decoderState.h +++ b/cpp/include/tensorrt_llm/runtime/decoderState.h @@ -66,7 +66,7 @@ class DecoderState WorldConfig const& worldConfig, BufferManager const& bufferManager); //! @brief Disable lookahead decoding. - void disableLookahead(RequestVector const& genRequests); + void disableLookahead(); //! @returns [batchSize], number of finished sequences per request, on gpu [[nodiscard]] TensorPtr getFinishedSum() const; diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 08cb4d407c17..5c35925af436 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -2604,9 +2604,15 @@ void TrtGptModelInflightBatching::changeSpecDecMode(ScheduledRequests const& sch mDecoderOutputBuffers.at(getFusedBufferId()).disableLookaheadDecoding(getMaxNumSequences()); mDecoder->disableLookahead( scheduledRequests.generationRequests, mDecoderInputBuffers.at(getFusedBufferId()).setupBatchSlots); - mDecoderState->disableLookahead(scheduledRequests.generationRequests); + mDecoderState->disableLookahead(); + for (auto const& llmReq : scheduledRequests.generationRequests) { + if (llmReq->mSeqSlot) + { + mDecoderState->setNumDecodingEngineTokens(llmReq->mSeqSlot.value(), 1); + } + if (llmReq->getNumDraftTokens() > 0) { llmReq->discardDraftTokens(llmReq->getNumDraftTokens()); diff --git a/cpp/tensorrt_llm/runtime/decoderState.cpp b/cpp/tensorrt_llm/runtime/decoderState.cpp index b5851dc1c2d2..d6fae46efdc4 100644 --- a/cpp/tensorrt_llm/runtime/decoderState.cpp +++ b/cpp/tensorrt_llm/runtime/decoderState.cpp @@ -15,7 +15,6 @@ */ #include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/runtime/runtimeKernels.h" @@ -444,7 +443,7 @@ void DecoderState::reshapeSpeculativeDecodingBuffers(SpeculativeDecodingMode con TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -void DecoderState::disableLookahead(RequestVector const& genRequests) +void DecoderState::disableLookahead() { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -457,14 +456,6 @@ void DecoderState::disableLookahead(RequestVector const& genRequests) auto const maxNewTokensShape = ITensor::makeShape({mMaxDecodingEngineTokens, mMaxNumSequences, mMaxBeamWidth}); mJointDecodingOutput->newTokensSteps->reshape(maxNewTokensShape); - for (auto const& llmReq : genRequests) - { - if (llmReq->mSeqSlot) - { - setNumDecodingEngineTokens(llmReq->mSeqSlot.value(), 1); - } - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } From f7ab50f2e7d8cc6e4c5437f3b9a04e3ce30afc1e Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Sun, 29 Jun 2025 14:37:45 +0000 Subject: [PATCH 2/4] refactor: Remove disableLookahead method from GptDecoderBatched and iGptDecoderBatched - Eliminated the disableLookahead method from both GptDecoderBatched and iGptDecoderBatched interfaces to streamline the API. - Updated TrtGptModelInflightBatching to utilize a new internal method for handling lookahead disabling, improving encapsulation and clarity. - Adjusted related logic in the setupDecoderStep method to maintain functionality without the removed method. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../tensorrt_llm/runtime/gptDecoderBatched.h | 14 +------ .../tensorrt_llm/runtime/iGptDecoderBatched.h | 10 ----- .../trtGptModelInflightBatching.cpp | 39 ++++++++++++++++++- .../trtGptModelInflightBatching.h | 2 + .../runtime/gptDecoderBatched.cpp | 33 ---------------- 5 files changed, 41 insertions(+), 57 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h index d0a9e726d13f..713fb23f5b32 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h @@ -26,12 +26,6 @@ #include "tensorrt_llm/runtime/worldConfig.h" #include -#include - -namespace tensorrt_llm::batch_manager -{ -class LlmRequest; -} // namespace tensorrt_llm::batch_manager namespace tensorrt_llm::runtime { @@ -41,8 +35,6 @@ class GptDecoderBatched : public IGptDecoderBatched { public: using CudaStreamPtr = std::shared_ptr; - using LlmRequestPtr = std::shared_ptr; - using RequestVector = std::vector; using TensorPtr = ITensor::SharedPtr; explicit GptDecoderBatched(CudaStreamPtr stream); @@ -50,8 +42,6 @@ class GptDecoderBatched : public IGptDecoderBatched void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) override; - void disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) override; - CudaEvent forwardAsync(decoder::DecoderState const& decoderState, decoder_batch::Input const& input) override; void forward(decoder::DecoderState const& decoderState, decoder_batch::Input const& input) override; @@ -60,12 +50,12 @@ class GptDecoderBatched : public IGptDecoderBatched [[nodiscard]] CudaEvent finalize(decoder::DecoderState const& decoderState, SizeType32 batchSlot, SamplingConfig const& samplingConfig, bool streaming) const override; - CudaStreamPtr getDecoderStream() const + [[nodiscard]] CudaStreamPtr getDecoderStream() const { return mDecoderStream; } - IGptDecoder& getUnderlyingDecoder() const + [[nodiscard]] IGptDecoder& getUnderlyingDecoder() const { return *mDecoder.get(); } diff --git a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h index 606ba3c98a48..17aacd851891 100644 --- a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h +++ b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h @@ -25,11 +25,6 @@ #include #include -namespace tensorrt_llm::batch_manager -{ -class LlmRequest; -} - namespace tensorrt_llm::runtime { class SamplingConfig; @@ -81,8 +76,6 @@ class IGptDecoderBatched { public: using CudaStreamPtr = std::shared_ptr; - using LlmRequestPtr = std::shared_ptr; - using RequestVector = std::vector; using TensorPtr = std::shared_ptr; //! @brief Setup the decoder before calling `forward()` @@ -90,9 +83,6 @@ class IGptDecoderBatched nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) = 0; - //! @brief Disable Lookahead decoding. - virtual void disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) = 0; - //! @brief Run one step for all requests without blocking the host process and return the token for synchronization. virtual CudaEvent forwardAsync(decoder::DecoderState const& decoderState, decoder_batch::Input const& input) = 0; diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 5c35925af436..808480956980 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -2515,6 +2515,42 @@ void TrtGptModelInflightBatching::changeBeamWidth(SizeType32 beamWidth) TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } +void TrtGptModelInflightBatching::disableLookaheadDecoder( + RequestVector const& genRequests, DecoderInputBuffers& inputBuffers) +{ + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); + + auto& batchSlots = inputBuffers.setupBatchSlots; + + std::vector samplingConfigs; + samplingConfigs.reserve(genRequests.size()); + auto batchSlotsRange = BufferRange(*batchSlots); + + SizeType32 batchIdx = 0; + for (auto const& llmReq : genRequests) + { + samplingConfigs.push_back(llmReq->mSamplingConfig); + batchSlotsRange[batchIdx] = llmReq->mSeqSlot.value(); + batchIdx += 1; + } + auto const batchSize = batchIdx; + std::optional samplingConfig; + if (batchSize > 0) + { + samplingConfig = SamplingConfig(samplingConfigs); + } + TensorPtr batchSlotsView = ITensor::slice(batchSlots, 0, batchSize); + + mDecoder->getUnderlyingDecoder().disableLookahead(samplingConfig, batchSize, batchSlots); + + auto const& decoderStream = mDecoder->getDecoderStream(); + CudaEvent event{}; + decoderStream->record(event); + mRuntime->getStreamPtr()->wait(event); + + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); +} + void TrtGptModelInflightBatching::changeSpecDecMode(ScheduledRequests const& scheduledRequests) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -2602,8 +2638,7 @@ void TrtGptModelInflightBatching::changeSpecDecMode(ScheduledRequests const& sch mDecodingConfig.setDecodingMode(executor::DecodingMode::Auto()); mBuffers.at(bufferId)->mLookaheadBuffers->disableLookaheadDecoding(); mDecoderOutputBuffers.at(getFusedBufferId()).disableLookaheadDecoding(getMaxNumSequences()); - mDecoder->disableLookahead( - scheduledRequests.generationRequests, mDecoderInputBuffers.at(getFusedBufferId()).setupBatchSlots); + disableLookaheadDecoder(scheduledRequests.generationRequests, mDecoderInputBuffers.at(getFusedBufferId())); mDecoderState->disableLookahead(); for (auto const& llmReq : scheduledRequests.generationRequests) diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h index 28d1767525cc..d76c415f9390 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h @@ -388,6 +388,8 @@ class TrtGptModelInflightBatching : public TrtGptModel /// @brief Change the speculative decoding mode. void changeSpecDecMode(ScheduledRequests const& scheduledRequests); + void disableLookaheadDecoder(RequestVector const& genRequests, DecoderInputBuffers& inputBuffers); + void prefetchNextPromptTableChunk(RequestVector const& contextRequests, bool isFirstChunk, SizeType32 bufferId); void remapInputTokensForPromptTable( diff --git a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp index 6df7b1634b88..d919cdc10672 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp @@ -19,8 +19,6 @@ #include "common.h" #include "decoderState.h" #include "iBuffer.h" -#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/kernels/decodingKernels.h" @@ -41,37 +39,6 @@ GptDecoderBatched::GptDecoderBatched(GptDecoderBatched::CudaStreamPtr stream) { } -void GptDecoderBatched::disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - std::vector samplingConfigs; - samplingConfigs.reserve(genRequests.size()); - auto batchSlotsRange = BufferRange(*batchSlots); - - SizeType32 batchIdx = 0; - for (auto const& llmReq : genRequests) - { - samplingConfigs.push_back(llmReq->mSamplingConfig); - batchSlotsRange[batchIdx] = llmReq->mSeqSlot.value(); - batchIdx += 1; - } - auto const batchSize = batchIdx; - std::optional samplingConfig; - if (batchSize > 0) - { - samplingConfig = SamplingConfig(samplingConfigs); - } - TensorPtr batchSlotsView = ITensor::slice(batchSlots, 0, batchSize); - mDecoder->disableLookahead(samplingConfig, batchSize, batchSlots); - - CudaEvent event{}; - mDecoderStream->record(event); - mRuntimeStream->wait(event); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - void GptDecoderBatched::setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) { From d279e0cec9164be55d245302e8eacca5bc9cb612 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Sun, 29 Jun 2025 14:50:32 +0000 Subject: [PATCH 3/4] refactor: Unify sampling configuration and batch slot handling in CreateNewDecoderRequests - Updated the operator() method to utilize std::optional for sampling configurations, improving flexibility in handling requests. - Introduced fillBatchSlots and fuseSamplingConfigs methods to streamline batch slot filling and sampling configuration fusion. - Refactored copySequenceLengths to return batch slots, enhancing clarity and separation of concerns in the code. - Use new methods to disable lookahead in trtGptModelInflightBatching. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.h | 6 +- .../createNewDecoderRequests.cpp | 67 ++++++++++++++----- .../trtGptModelInflightBatching.cpp | 33 ++------- cpp/tensorrt_llm/nanobind/bindings.cpp | 12 ---- .../nanobind/common/customCasters.h | 2 - cpp/tensorrt_llm/pybind/bindings.cpp | 10 --- tensorrt_llm/_torch/pyexecutor/sampler.py | 6 +- 7 files changed, 63 insertions(+), 73 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index 0978905b5e2f..73b0018ce0de 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -70,7 +70,7 @@ class CreateNewDecoderRequests : Algorithm { } - std::tuple, std::vector, + std::tuple, std::vector, std::vector> operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, @@ -86,6 +86,10 @@ class CreateNewDecoderRequests : Algorithm runtime::CudaStream const& runtimeStream, runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const; + static TensorPtr fillBatchSlots(RequestVector const& requests, DecoderInputBuffers& inputBuffers); + + static std::optional fuseSamplingConfigs(RequestVector const& requests); + private: bool mSpeculativeDecodingFastLogits; bool mIsLeaderInOrchMode; diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 3335d69a015b..d7cb9d106681 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -54,7 +54,8 @@ using OptionalRef = tensorrt_llm::common::OptionalRef; namespace { -void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffers& inputBuffers, +//! @brief Fills the seqSlots and sequence lengths in the inputBuffers. +TensorPtr copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffers& inputBuffers, ITensor& sequenceLengths, SizeType32 beamWidth, runtime::CudaStream const& stream) { auto const bufferManager = BufferManager{std::make_shared(stream.get())}; @@ -82,6 +83,7 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe } // copy sequence lengths + if (!contextRequests.empty()) { auto batchSlotsDeviceView = tr::ITensor::slice(inputBuffers.setupBatchSlotsDevice, 0, batchSize); auto fillValuesViewDevice = tr::ITensor::slice(inputBuffers.fillValuesDevice, 0, batchSize); @@ -90,6 +92,8 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe bufferManager.copy(*fillValuesView, *fillValuesViewDevice); tr::kernels::invokeFillBatch(sequenceLengths, *batchSlotsDeviceView, beamWidth, *fillValuesViewDevice, stream); } + + return batchSlotsView; } /// @brief Retrieve the embedding bias from the request. This potentially makes a copy of the tensor @@ -131,7 +135,46 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe } // namespace -std::tuple, std::vector, +// Similar to copySequenceLengths, but only fills the seqSlots. +TensorPtr CreateNewDecoderRequests::fillBatchSlots(RequestVector const& requests, DecoderInputBuffers& inputBuffers) +{ + auto const batchSize = requests.size(); + auto batchSlotsView = tr::ITensor::slice(inputBuffers.setupBatchSlots, 0, batchSize); + + auto batchSlotsRange = tr::BufferRange(*batchSlotsView); + + // fill buffers on host + SizeType32 batchIdx{0}; + for (auto const& llmReq : requests) + { + auto const seqSlot = llmReq->mSeqSlot.value(); + batchSlotsRange[batchIdx] = seqSlot; + ++batchIdx; + } + + // TODO: copy to device and use in GptDecoder + // manager.copy(*batchSlotsView, *batchSlotsDeviceView); + + return batchSlotsView; +} + +std::optional CreateNewDecoderRequests::fuseSamplingConfigs(RequestVector const& requests) +{ + if (requests.empty()) + { + return std::nullopt; + } + + std::vector samplingConfigs; + samplingConfigs.reserve(requests.size()); + for (auto const& llmReq : requests) + { + samplingConfigs.push_back(llmReq->mSamplingConfig); + } + return SamplingConfig(samplingConfigs); +} + +std::tuple, std::vector, std::vector> CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, nvinfer1::DataType logitsType, @@ -146,29 +189,17 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru std::copy_if(contextRequests.begin(), contextRequests.end(), std::back_inserter(finishedContextRequests), [](auto const& llmReq) { return llmReq->isLastContextChunk(); }); - if (!finishedContextRequests.empty()) - { - copySequenceLengths( - finishedContextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, runtimeStream); - } + auto batchSlotsView = copySequenceLengths( + finishedContextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, runtimeStream); auto [lookaheadPrompt, lookaheadAlgoConfigs] = createDecoderRequests(finishedContextRequests, inputBuffers.inputsIds, decodingConfig, decoderState, logitsType, modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, medusaBuffers); - auto const batchSize = finishedContextRequests.size(); - - std::vector samplingConfigs; - samplingConfigs.reserve(batchSize); - for (auto const& llmReq : finishedContextRequests) - { - samplingConfigs.push_back(llmReq->mSamplingConfig); - } - - TensorPtr batchSlotsView = runtime::ITensor::slice(inputBuffers.setupBatchSlots, 0, batchSize); + auto samplingConfig = fuseSamplingConfigs(finishedContextRequests); TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {std::move(batchSlotsView), std::move(samplingConfigs), std::move(lookaheadPrompt), + return {std::move(batchSlotsView), std::move(samplingConfig), std::move(lookaheadPrompt), std::move(lookaheadAlgoConfigs)}; } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 808480956980..7114e399bdd4 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -1865,7 +1865,7 @@ void TrtGptModelInflightBatching::setupDecoderStep( { auto const logitsType = mRuntime->getEngine().getTensorDataType("logits"); - auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] + auto [batchSlots, samplingConfig, lookaheadPrompt, lookaheadAlgoConfigs] = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, contextRequests, logitsType, inputBuffers, *mDecoderState, mRuntime->getStream(), *mDecoder->getDecoderStream(), getMaxSequenceLen(), mOperatingBeamWidth, buffers.mMedusaBuffers); @@ -1873,14 +1873,13 @@ void TrtGptModelInflightBatching::setupDecoderStep( auto const localBatchSize = batchSlots->getSize(); if (localBatchSize > 0) { - auto samplingConfig = SamplingConfig(samplingConfigs); - mDecoder->getUnderlyingDecoder().setup(samplingConfig, localBatchSize, batchSlots, + mDecoder->getUnderlyingDecoder().setup(samplingConfig.value(), localBatchSize, batchSlots, {mDecoderState->getJointDecodingOutput()}, mModelConfig.getDataType(), lookaheadPrompt, lookaheadAlgoConfigs); - auto const& stream = mDecoder->getDecoderStream(); + auto const& decoderStream = mDecoder->getDecoderStream(); CudaEvent event{}; - stream->record(event); + decoderStream->record(event); mRuntime->getStreamPtr()->wait(event); } } @@ -2520,28 +2519,10 @@ void TrtGptModelInflightBatching::disableLookaheadDecoder( { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto& batchSlots = inputBuffers.setupBatchSlots; + auto batchSlots = CreateNewDecoderRequests::fillBatchSlots(genRequests, inputBuffers); + auto samplingConfig = CreateNewDecoderRequests::fuseSamplingConfigs(genRequests); - std::vector samplingConfigs; - samplingConfigs.reserve(genRequests.size()); - auto batchSlotsRange = BufferRange(*batchSlots); - - SizeType32 batchIdx = 0; - for (auto const& llmReq : genRequests) - { - samplingConfigs.push_back(llmReq->mSamplingConfig); - batchSlotsRange[batchIdx] = llmReq->mSeqSlot.value(); - batchIdx += 1; - } - auto const batchSize = batchIdx; - std::optional samplingConfig; - if (batchSize > 0) - { - samplingConfig = SamplingConfig(samplingConfigs); - } - TensorPtr batchSlotsView = ITensor::slice(batchSlots, 0, batchSize); - - mDecoder->getUnderlyingDecoder().disableLookahead(samplingConfig, batchSize, batchSlots); + mDecoder->getUnderlyingDecoder().disableLookahead(samplingConfig, batchSlots->getSize(), batchSlots); auto const& decoderStream = mDecoder->getDecoderStream(); CudaEvent event{}; diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index c951f967c205..ed0e0a214dd4 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -64,14 +64,6 @@ using OptVec = std::optional>; #error "TRTLLM_NB_MODULE must be defined" #endif -namespace -{ -tr::SamplingConfig makeSamplingConfig(std::vector const& configs) -{ - return tr::SamplingConfig(configs); -} -} // namespace - NB_MODULE(TRTLLM_NB_MODULE, m) { m.doc() = "TensorRT-LLM Python bindings for C++ runtime"; @@ -425,10 +417,6 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def("__setstate__", SamplingConfigSetState) .def("__eq__", &tr::SamplingConfig::operator==); - nb::bind_vector>(m, "SamplingConfigVector"); - - m.def("make_sampling_config", &makeSamplingConfig, nb::arg("configs")); - nb::class_(m, "GptJsonConfig") .def(nb::init>(), diff --git a/cpp/tensorrt_llm/nanobind/common/customCasters.h b/cpp/tensorrt_llm/nanobind/common/customCasters.h index 432ce5c26b73..25d20b54cbea 100644 --- a/cpp/tensorrt_llm/nanobind/common/customCasters.h +++ b/cpp/tensorrt_llm/nanobind/common/customCasters.h @@ -21,7 +21,6 @@ #include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/common/optionalRef.h" #include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/samplingConfig.h" #include "tensorrt_llm/runtime/torch.h" #include "tensorrt_llm/runtime/torchView.h" @@ -45,7 +44,6 @@ // Opaque bindings NB_MAKE_OPAQUE(tensorrt_llm::batch_manager::ReqIdsSet) NB_MAKE_OPAQUE(std::vector) -NB_MAKE_OPAQUE(std::vector) namespace nb = nanobind; diff --git a/cpp/tensorrt_llm/pybind/bindings.cpp b/cpp/tensorrt_llm/pybind/bindings.cpp index cdc9736db096..4260694b116d 100644 --- a/cpp/tensorrt_llm/pybind/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/bindings.cpp @@ -58,14 +58,6 @@ using OptVec = std::optional>; #error "TRTLLM_PYBIND_MODULE must be defined" #endif -namespace -{ -tr::SamplingConfig makeSamplingConfig(std::vector const& configs) -{ - return tr::SamplingConfig(configs); -} -} // namespace - PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m) { m.doc() = "TensorRT-LLM Python bindings for C++ runtime"; @@ -415,8 +407,6 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m) .def(py::pickle(SamplingConfigGetState, SamplingConfigSetState)) .def("__eq__", &tr::SamplingConfig::operator==); - m.def("make_sampling_config", &makeSamplingConfig, py::arg("configs")); - py::class_(m, "GptJsonConfig") .def(py::init>(), diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index e6d19a9df46c..d830d63f8954 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -8,8 +8,7 @@ from tensorrt_llm._torch.pyexecutor.make_decoding_batch_input_output import \ MakeDecodingBatchInputOutput from tensorrt_llm._utils import nvtx_range, torch_dtype_to_binding -from tensorrt_llm.bindings import (CudaStream, DataType, ModelConfig, - WorldConfig, make_sampling_config) +from tensorrt_llm.bindings import CudaStream, DataType, ModelConfig, WorldConfig from tensorrt_llm.bindings.executor import (DecodingConfig, DecodingMode, ExecutorConfig, FinishReason) from tensorrt_llm.bindings.internal.algorithms import CreateNewDecoderRequests @@ -885,7 +884,7 @@ def _instantiate_algorithms(self): @torch.inference_mode() @nvtx_range("setup_sampler_step") def setup_sampler_step(self, requests): - batch_slots, sampling_configs, lookahead_prompt, lookahead_algo_configs = self.algs.create_new_decoder_requests( + batch_slots, sampling_config, lookahead_prompt, lookahead_algo_configs = self.algs.create_new_decoder_requests( self.model_config, self.world_config, self.decoding_config, requests.context_requests, self.logits_datatype, self.store["decoder_input_buffers"][self.micro_batch_idx], @@ -895,7 +894,6 @@ def setup_sampler_step(self, requests): local_batch_size = len(batch_slots) if local_batch_size > 0: - sampling_config = make_sampling_config(sampling_configs) self.algs.decoder.underlying_decoder().setup( sampling_config, local_batch_size, batch_slots, self.store["decoder_state"].joint_decoding_output, From 6062c69f6eb859019eba7efce8b5b63cd694470e Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 09:53:38 +0000 Subject: [PATCH 4/4] refactor: Update CreateNewDecoderRequests to handle context requests more effectively - Filter out finished context requests from the contextRequests vector before calling CreateNewDecoderRequests. - Modified the operator() method to directly use finishedContextRequests, enhancing clarity. - Updated the handling of sampling configurations to ensure a valid value is returned, improving robustness. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.h | 2 +- .../createNewDecoderRequests.cpp | 18 +++++++------- .../trtGptModelInflightBatching.cpp | 24 ++++++++++++------- tensorrt_llm/_torch/pyexecutor/sampler.py | 24 ++++++++++++------- 4 files changed, 41 insertions(+), 27 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index 73b0018ce0de..0526674df373 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -70,7 +70,7 @@ class CreateNewDecoderRequests : Algorithm { } - std::tuple, std::vector, + std::tuple, std::vector> operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index d7cb9d106681..3157f3d708e8 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -174,7 +174,7 @@ std::optional CreateNewDecoderRequests::fuseSamplingConfigs(Requ return SamplingConfig(samplingConfigs); } -std::tuple, std::vector, +std::tuple, std::vector> CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, nvinfer1::DataType logitsType, @@ -185,21 +185,21 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); NVTX3_SCOPED_RANGE(CreateNewDecoderRequests); - RequestVector finishedContextRequests; - std::copy_if(contextRequests.begin(), contextRequests.end(), std::back_inserter(finishedContextRequests), - [](auto const& llmReq) { return llmReq->isLastContextChunk(); }); + TLLM_CHECK_WITH_INFO( + !contextRequests.empty(), "CreateNewDecoderRequests should be called with at least one request"); auto batchSlotsView = copySequenceLengths( - finishedContextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, runtimeStream); + contextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, runtimeStream); auto [lookaheadPrompt, lookaheadAlgoConfigs] - = createDecoderRequests(finishedContextRequests, inputBuffers.inputsIds, decodingConfig, decoderState, - logitsType, modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, medusaBuffers); + = createDecoderRequests(contextRequests, inputBuffers.inputsIds, decodingConfig, decoderState, logitsType, + modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, medusaBuffers); - auto samplingConfig = fuseSamplingConfigs(finishedContextRequests); + auto samplingConfig = fuseSamplingConfigs(contextRequests); + TLLM_CHECK(samplingConfig.has_value()); TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {std::move(batchSlotsView), std::move(samplingConfig), std::move(lookaheadPrompt), + return {std::move(batchSlotsView), std::move(samplingConfig.value()), std::move(lookaheadPrompt), std::move(lookaheadAlgoConfigs)}; } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 7114e399bdd4..9f0b4903dd6d 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -1863,17 +1863,23 @@ void TrtGptModelInflightBatching::setupDecoderStep( if (mWorldConfig.isLastPipelineParallelRank() && !contextRequests.empty()) { - auto const logitsType = mRuntime->getEngine().getTensorDataType("logits"); + RequestVector finishedContextRequests; + std::copy_if(contextRequests.begin(), contextRequests.end(), std::back_inserter(finishedContextRequests), + [](auto const& llmReq) { return llmReq->isLastContextChunk(); }); - auto [batchSlots, samplingConfig, lookaheadPrompt, lookaheadAlgoConfigs] - = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, contextRequests, logitsType, - inputBuffers, *mDecoderState, mRuntime->getStream(), *mDecoder->getDecoderStream(), getMaxSequenceLen(), - mOperatingBeamWidth, buffers.mMedusaBuffers); - - auto const localBatchSize = batchSlots->getSize(); - if (localBatchSize > 0) + if (!finishedContextRequests.empty()) { - mDecoder->getUnderlyingDecoder().setup(samplingConfig.value(), localBatchSize, batchSlots, + auto const logitsType = mRuntime->getEngine().getTensorDataType("logits"); + + auto [batchSlots, samplingConfig, lookaheadPrompt, lookaheadAlgoConfigs] + = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, finishedContextRequests, + logitsType, inputBuffers, *mDecoderState, mRuntime->getStream(), *mDecoder->getDecoderStream(), + getMaxSequenceLen(), mOperatingBeamWidth, buffers.mMedusaBuffers); + + auto const localBatchSize = batchSlots->getSize(); + TLLM_CHECK_WITH_INFO(localBatchSize > 0, "Decoder setup should be called with at least one request"); + + mDecoder->getUnderlyingDecoder().setup(samplingConfig, localBatchSize, batchSlots, {mDecoderState->getJointDecodingOutput()}, mModelConfig.getDataType(), lookaheadPrompt, lookaheadAlgoConfigs); diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index d830d63f8954..99ea5f5b6970 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -884,21 +884,29 @@ def _instantiate_algorithms(self): @torch.inference_mode() @nvtx_range("setup_sampler_step") def setup_sampler_step(self, requests): + finished_context_requests = [ + r for r in requests if r.is_last_context_chunk + ] + + if len(finished_context_requests) == 0: + return + batch_slots, sampling_config, lookahead_prompt, lookahead_algo_configs = self.algs.create_new_decoder_requests( self.model_config, self.world_config, self.decoding_config, - requests.context_requests, self.logits_datatype, + finished_context_requests, self.logits_datatype, self.store["decoder_input_buffers"][self.micro_batch_idx], self.store["decoder_state"], self.store["cuda_stream"], self.algs.decoder.decoder_stream, self.executor_config.max_seq_len, - self.beam_width(requests.context_requests)) + self.beam_width(finished_context_requests)) local_batch_size = len(batch_slots) - if local_batch_size > 0: - self.algs.decoder.underlying_decoder().setup( - sampling_config, local_batch_size, batch_slots, - self.store["decoder_state"].joint_decoding_output, - self.model_config.data_type, lookahead_prompt, - lookahead_algo_configs) + assert local_batch_size > 0, "Decoder setup should be called with at least one request" + + self.algs.decoder.underlying_decoder().setup( + sampling_config, local_batch_size, batch_slots, + self.store["decoder_state"].joint_decoding_output, + self.model_config.data_type, lookahead_prompt, + lookahead_algo_configs) adp = [ r for r in requests.generation_requests if r.is_attention_dp_dummy