From aae96888e42e371d07c99184c4d805b343093bef Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 11:15:23 +0000 Subject: [PATCH 01/17] refactor: Minor cleanup in CreateNewDecoderRequests Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../tensorrt_llm/batch_manager/createNewDecoderRequests.h | 3 +-- cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index ce42493879e9..606a8017f8ca 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -89,8 +89,7 @@ class CreateNewDecoderRequests : Algorithm SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const; private: - //! @brief Initialize the decoder at `batchSlot` with a new `request`. Exposed only for static batching via - //! GptDecoderBatched::newBatch() + //! @brief Initialize the decoder at `batchSlot` with a new `request`. static void newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 84fd13d5c18a..0d19a1588d61 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -44,6 +44,7 @@ namespace tensorrt_llm::batch_manager using SizeType32 = CreateNewDecoderRequests::SizeType32; using TensorPtr = CreateNewDecoderRequests::TensorPtr; +using SharedConstPtr = CreateNewDecoderRequests::SharedConstPtr; namespace { From 7f10367f4d52a7a0bd3e8619f073d1cc61567871 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 12:58:17 +0000 Subject: [PATCH 02/17] refactor: remove BufferManager from CreateNewDecoderRequests parameters - Updated CreateNewDecoderRequests class and related methods to remove BufferManager as a parameter. - CreateNewDecoderRequests already has runtimeStream as a parameter, so we can use it to create the BufferManager. - This makes it easier to differentiate between the runtime stream and the decoder stream. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.h | 9 +++-- .../createNewDecoderRequests.cpp | 34 ++++++++++--------- .../trtGptModelInflightBatching.cpp | 6 ++-- .../nanobind/batch_manager/algorithms.cpp | 16 ++++----- .../pybind/batch_manager/algorithms.cpp | 16 ++++----- cpp/tests/runtime/gptDecoderBatchedTest.cpp | 7 ++-- tensorrt_llm/_torch/pyexecutor/sampler.py | 3 +- 7 files changed, 43 insertions(+), 48 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index 606a8017f8ca..44176677eb32 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -75,16 +75,15 @@ class CreateNewDecoderRequests : Algorithm std::vector> operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - runtime::BufferManager const& bufferManager, nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, - runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, - SizeType32 maxSequenceLength, SizeType32 beamWidth, OptionalRef medusaBuffers) const; + nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, + CudaStream const& runtimeStream, CudaStream const& decoderStream, SizeType32 maxSequenceLength, + SizeType32 beamWidth, OptionalRef medusaBuffers) const; [[nodiscard]] std::tuple, std::vector> createDecoderRequests(RequestVector const& finishedContextRequests, TensorPtr const& inputIds, executor::DecodingConfig const& decodingConfig, runtime::decoder::DecoderState& decoderState, - runtime::BufferManager const& bufferManager, nvinfer1::DataType logitsType, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, + nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, runtime::CudaStream const& runtimeStream, runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const; diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 0d19a1588d61..821dd9d97e22 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -50,9 +50,10 @@ namespace { void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffers& inputBuffers, - ITensor& sequenceLengths, SizeType32 beamWidth, runtime::BufferManager const& manager, - runtime::CudaStream const& stream) + ITensor& sequenceLengths, SizeType32 beamWidth, runtime::CudaStream const& stream) { + auto const bufferManager = BufferManager{std::make_shared(stream.get())}; + auto const batchSize = contextRequests.size(); auto batchSlotsView = tr::ITensor::slice(inputBuffers.setupBatchSlots, 0, batchSize); auto fillValuesView = tr::ITensor::slice(inputBuffers.fillValues, 0, batchSize); @@ -80,8 +81,8 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe auto batchSlotsDeviceView = tr::ITensor::slice(inputBuffers.setupBatchSlotsDevice, 0, batchSize); auto fillValuesViewDevice = tr::ITensor::slice(inputBuffers.fillValuesDevice, 0, batchSize); - manager.copy(*batchSlotsView, *batchSlotsDeviceView); - manager.copy(*fillValuesView, *fillValuesViewDevice); + bufferManager.copy(*batchSlotsView, *batchSlotsDeviceView); + bufferManager.copy(*fillValuesView, *fillValuesViewDevice); tr::kernels::invokeFillBatch(sequenceLengths, *batchSlotsDeviceView, beamWidth, *fillValuesViewDevice, stream); } } @@ -128,10 +129,10 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe std::tuple, std::vector, std::vector> CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - runtime::BufferManager const& bufferManager, nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, - runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, - SizeType32 maxSequenceLength, SizeType32 beamWidth, OptionalRef medusaBuffers) const + executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, nvinfer1::DataType logitsType, + DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, + CudaStream const& decoderStream, SizeType32 maxSequenceLength, SizeType32 beamWidth, + OptionalRef medusaBuffers) const { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); NVTX3_SCOPED_RANGE(CreateNewDecoderRequests); @@ -142,13 +143,13 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru if (!finishedContextRequests.empty()) { - copySequenceLengths(finishedContextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, - bufferManager, runtimeStream); + copySequenceLengths( + finishedContextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, runtimeStream); } - auto [lookaheadPrompt, lookaheadAlgoConfigs] = createDecoderRequests(finishedContextRequests, - inputBuffers.inputsIds, decodingConfig, decoderState, bufferManager, logitsType, modelConfig, worldConfig, - runtimeStream, decoderStream, maxSequenceLength, medusaBuffers); + auto [lookaheadPrompt, lookaheadAlgoConfigs] + = createDecoderRequests(finishedContextRequests, inputBuffers.inputsIds, decodingConfig, decoderState, + logitsType, modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, medusaBuffers); auto const batchSize = finishedContextRequests.size(); @@ -558,11 +559,12 @@ void CreateNewDecoderRequests::newRequestEagle(SizeType32 batchIdx, runtime::dec std::tuple, std::vector> CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedContextRequests, TensorPtr const& inputIds, executor::DecodingConfig const& decodingConfig, runtime::decoder::DecoderState& decoderState, - BufferManager const& bufferManager, nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, runtime::CudaStream const& runtimeStream, - runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, + nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, + runtime::CudaStream const& runtimeStream, runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const { + auto const bufferManager = BufferManager{std::make_shared(runtimeStream.get())}; + unsigned decoderInputSize{0}; for (auto const& llmReq : finishedContextRequests) { diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index d42d798f68ba..08cb4d407c17 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -1866,9 +1866,9 @@ void TrtGptModelInflightBatching::setupDecoderStep( auto const logitsType = mRuntime->getEngine().getTensorDataType("logits"); auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] - = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, contextRequests, - mRuntime->getBufferManager(), logitsType, inputBuffers, *mDecoderState, mRuntime->getStream(), - *mDecoder->getDecoderStream(), getMaxSequenceLen(), mOperatingBeamWidth, buffers.mMedusaBuffers); + = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, contextRequests, logitsType, + inputBuffers, *mDecoderState, mRuntime->getStream(), *mDecoder->getDecoderStream(), getMaxSequenceLen(), + mOperatingBeamWidth, buffers.mMedusaBuffers); auto const localBatchSize = batchSlots->getSize(); if (localBatchSize > 0) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp index 357a502bcce6..1944784eee2e 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp @@ -103,23 +103,21 @@ void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_ "__call__", [](CreateNewDecoderRequests& self, tr::ModelConfig const& modelConfig, tr::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - tr::BufferManager const& bufferManager, nvinfer1::DataType logitsType, - DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, - tensorrt_llm::runtime::CudaStream const& runtimeStream, + nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, + runtime::decoder::DecoderState& decoderState, tensorrt_llm::runtime::CudaStream const& runtimeStream, tensorrt_llm::runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, SizeType32 beamWidth) { OptionalRef medusaBuffers = std::nullopt; - auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] = self(modelConfig, - worldConfig, decodingConfig, contextRequests, bufferManager, logitsType, inputBuffers, decoderState, - runtimeStream, decoderStream, maxSequenceLength, beamWidth, medusaBuffers); + auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] + = self(modelConfig, worldConfig, decodingConfig, contextRequests, logitsType, inputBuffers, + decoderState, runtimeStream, decoderStream, maxSequenceLength, beamWidth, medusaBuffers); return std::tuple{runtime::Torch::tensor(batchSlots), std::move(samplingConfigs), std::move(lookaheadPrompt), std::move(lookaheadAlgoConfigs)}; }, nb::arg("model_config"), nb::arg("world_config"), nb::arg("decoding_config"), nb::arg("context_requests"), - nb::arg("buffer_manager"), nb::arg("logits_type"), nb::arg("decoder_input_buffers"), - nb::arg("decoder_state"), nb::arg("runtime_stream"), nb::arg("decoder_stream"), - nb::arg("max_sequence_length"), nb::arg("beam_width")) + nb::arg("logits_type"), nb::arg("decoder_input_buffers"), nb::arg("decoder_state"), + nb::arg("runtime_stream"), nb::arg("decoder_stream"), nb::arg("max_sequence_length"), nb::arg("beam_width")) .def("name", [](CreateNewDecoderRequests const&) { return CreateNewDecoderRequests::name; }); } diff --git a/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp index f098398b6244..8f0cc3315cde 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp @@ -105,23 +105,21 @@ void tensorrt_llm::pybind::batch_manager::algorithms::initBindings(pybind11::mod "__call__", [](CreateNewDecoderRequests& self, tr::ModelConfig const& modelConfig, tr::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - tr::BufferManager const& bufferManager, nvinfer1::DataType logitsType, - DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, - tensorrt_llm::runtime::CudaStream const& runtimeStream, + nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, + runtime::decoder::DecoderState& decoderState, tensorrt_llm::runtime::CudaStream const& runtimeStream, tensorrt_llm::runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, SizeType32 beamWidth) { OptionalRef medusaBuffers = std::nullopt; - auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] = self(modelConfig, - worldConfig, decodingConfig, contextRequests, bufferManager, logitsType, inputBuffers, decoderState, - runtimeStream, decoderStream, maxSequenceLength, beamWidth, medusaBuffers); + auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] + = self(modelConfig, worldConfig, decodingConfig, contextRequests, logitsType, inputBuffers, + decoderState, runtimeStream, decoderStream, maxSequenceLength, beamWidth, medusaBuffers); return std::tuple{runtime::Torch::tensor(batchSlots), std::move(samplingConfigs), std::move(lookaheadPrompt), std::move(lookaheadAlgoConfigs)}; }, py::arg("model_config"), py::arg("world_config"), py::arg("decoding_config"), py::arg("context_requests"), - py::arg("buffer_manager"), py::arg("logits_type"), py::arg("decoder_input_buffers"), - py::arg("decoder_state"), py::arg("runtime_stream"), py::arg("decoder_stream"), - py::arg("max_sequence_length"), py::arg("beam_width")) + py::arg("logits_type"), py::arg("decoder_input_buffers"), py::arg("decoder_state"), + py::arg("runtime_stream"), py::arg("decoder_stream"), py::arg("max_sequence_length"), py::arg("beam_width")) .def("name", [](CreateNewDecoderRequests const&) { return CreateNewDecoderRequests::name; }); } diff --git a/cpp/tests/runtime/gptDecoderBatchedTest.cpp b/cpp/tests/runtime/gptDecoderBatchedTest.cpp index 7c152f48a9e8..338b974aa0d8 100644 --- a/cpp/tests/runtime/gptDecoderBatchedTest.cpp +++ b/cpp/tests/runtime/gptDecoderBatchedTest.cpp @@ -104,15 +104,14 @@ void newRequests(std::vector> const& requests, T SizeType32 maxSequenceLength, tb::DecoderInputBuffers& inputBuffers, decoder::DecoderState& decoderState) { auto const& decoderStream = *decoder.getDecoderStream(); - auto const bufferManager = BufferManager{std::make_shared(runtimeStream.get())}; auto batchSlotsRange = BufferRange(*batchSlots); auto const localBatchSize = batchSlots->getSize(); tb::CreateNewDecoderRequests createNewDecoderRequests(false, false, false); - auto [lookaheadPrompt, lookaheadAlgoConfigs] = createNewDecoderRequests.createDecoderRequests(requests, - inputBuffers.inputsIds, decodingConfig, decoderState, bufferManager, logitsType, modelConfig, worldConfig, - runtimeStream, decoderStream, maxSequenceLength, std::nullopt); + auto [lookaheadPrompt, lookaheadAlgoConfigs] + = createNewDecoderRequests.createDecoderRequests(requests, inputBuffers.inputsIds, decodingConfig, decoderState, + logitsType, modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, std::nullopt); std::vector samplingConfigs; samplingConfigs.reserve(requests.size()); diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 73e0acf03f8c..e17c029e703c 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -749,8 +749,7 @@ def _instantiate_algorithms(self): def setup_sampler_step(self, requests): batch_slots, sampling_configs, lookahead_prompt, lookahead_algo_configs = self.algs.create_new_decoder_requests( self.model_config, self.world_config, self.decoding_config, - requests.context_requests, self.store["buffer_manager"], - self.logits_datatype, + requests.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, From 92aabea0c798777c498a05a8efc05418f5db3760 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 11:16:17 +0000 Subject: [PATCH 03/17] refactor: Move embedding bias initialization to a separate function in CreateNewDecoderRequests Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 39 ++++++++++++------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 821dd9d97e22..4065c1fe5ab2 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -167,6 +167,30 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru std::move(lookaheadAlgoConfigs)}; } +namespace +{ + +void initializeEmbeddingBias(DecodingInput& dJointInput, TensorPtr const& embeddingBiasTensor, SizeType32 batchSlot, + runtime::ModelConfig const& modelConfig, BufferManager const& manager) +{ + TensorPtr const embeddingBiasSlice = ITensor::slice(constPointerCast(dJointInput.embeddingBias), batchSlot, 1); + if (embeddingBiasTensor) + { + TLLM_CHECK(embeddingBiasTensor->getShape().nbDims == 2); + TLLM_CHECK(embeddingBiasTensor->getShape().d[0] == 1); + TLLM_CHECK_WITH_INFO(embeddingBiasTensor->getShape().d[1] == modelConfig.getVocabSize(), + "The embedding bias shape is not as expected. Expected last dimension to be same as vocab size: %d.", + modelConfig.getVocabSize()); + manager.copy(*embeddingBiasTensor, *embeddingBiasSlice); + } + else + { + manager.setZero(*embeddingBiasSlice); + } +} + +} // namespace + void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, @@ -208,20 +232,7 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; runtime::kernels::invokeFill(*endIdTensorPtr, endId, decoderStream); - TensorPtr const embeddingBiasSlice = ITensor::slice(constPointerCast(dJointInput.embeddingBias), batchSlot, 1); - if (request.embeddingBias) - { - TLLM_CHECK(request.embeddingBias->getShape().nbDims == 2); - TLLM_CHECK(request.embeddingBias->getShape().d[0] == 1); - TLLM_CHECK_WITH_INFO(request.embeddingBias->getShape().d[1] == modelConfig.getVocabSize(), - "The embedding bias shape is not as expected. Expected last dimension to be same as vocab size: %d.", - modelConfig.getVocabSize()); - manager.copy(*request.embeddingBias, *embeddingBiasSlice); - } - else - { - manager.setZero(*embeddingBiasSlice); - } + initializeEmbeddingBias(dJointInput, request.embeddingBias, batchSlot, modelConfig, manager); auto setupWords = [](std::vector& jointWordsLists, TensorPtr const& requestWordsList, SharedConstPtr& jointWordsPtrs, SharedConstPtr& jointWordsLens, SizeType32& jointMaxWordsLen, From a70dec7384cc2d09bfc9b1e1d1a0263d1c0aa111 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 11:22:02 +0000 Subject: [PATCH 04/17] refactor: Extract setupWords function in CreateNewDecoderRequests - Moved the setupWords logic from a lambda expression into a dedicated function, improving code clarity and maintainability. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 45 +++++++++---------- 1 file changed, 22 insertions(+), 23 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 4065c1fe5ab2..7f4db81bf446 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -189,6 +189,28 @@ void initializeEmbeddingBias(DecodingInput& dJointInput, TensorPtr const& embedd } } +void setupWords(std::vector& jointWordsLists, TensorPtr const& requestWordsList, + SharedConstPtr& jointWordsPtrs, SharedConstPtr& jointWordsLens, SizeType32& jointMaxWordsLen, SizeType32 batchSlot) +{ + if (requestWordsList) + { + auto const wordsLen = requestWordsList->getShape().d[1]; + BufferRange(*constPointerCast(jointWordsPtrs))[batchSlot] + = runtime::bufferCast(*requestWordsList); + runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = wordsLen; + // FIXME: this is monotonically growing size + jointMaxWordsLen = std::max(static_cast(wordsLen), jointMaxWordsLen); + + // NOTE: jointWordsList is not used in gptDecoder, but required to keep WordsList's + // memory allocated + jointWordsLists[batchSlot] = requestWordsList; + } + else + { + runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = 0; + } +}; + } // namespace void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, @@ -234,29 +256,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder initializeEmbeddingBias(dJointInput, request.embeddingBias, batchSlot, modelConfig, manager); - auto setupWords = [](std::vector& jointWordsLists, TensorPtr const& requestWordsList, - SharedConstPtr& jointWordsPtrs, SharedConstPtr& jointWordsLens, SizeType32& jointMaxWordsLen, - SizeType32 batchSlot) - { - if (requestWordsList) - { - auto const wordsLen = requestWordsList->getShape().d[1]; - BufferRange(*constPointerCast(jointWordsPtrs))[batchSlot] - = runtime::bufferCast(*requestWordsList); - runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = wordsLen; - // FIXME: this is monotonically growing size - jointMaxWordsLen = std::max(static_cast(wordsLen), jointMaxWordsLen); - - // NOTE: jointWordsList is not used in gptDecoder, but required to keep WordsList's - // memory allocated - jointWordsLists[batchSlot] = requestWordsList; - } - else - { - runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = 0; - } - }; - setupWords(dJointInput.stopWordsLists, request.stopWordsList, dJointInput.stopWordsPtrs, dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot); From 8fc612d483584a0492ef0b7c7f339963fd23b735 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 11:29:20 +0000 Subject: [PATCH 05/17] refactor: Move newRequestSpeculativeDecoding out of newRequest - Moved the setting of numDecodingEngineTokens to the newRequest function for better clarity. - Reorganized speculative execution logic to improve readability and maintainability. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.h | 3 +- .../createNewDecoderRequests.cpp | 36 +++++++++++-------- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index 44176677eb32..c35ba95d5a0a 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -91,8 +91,7 @@ class CreateNewDecoderRequests : Algorithm //! @brief Initialize the decoder at `batchSlot` with a new `request`. static void newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, - runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, - SizeType32 maxSequenceLength); + runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream, SizeType32 maxSequenceLength); //! @brief Setups decoder internal tensors for new speculative decoding request static void newRequestSpeculativeDecoding(SizeType32 batchIdx, runtime::decoder_batch::Request const& request, diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 7f4db81bf446..7964a279d3c7 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -215,8 +215,7 @@ void setupWords(std::vector& jointWordsLists, Tenso void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, - runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, - SizeType32 maxSequenceLength) + runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream, SizeType32 maxSequenceLength) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -249,7 +248,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder auto& dJointInput = decoderState.getJointDecodingInput(); dJointInput.beamWidths.at(batchSlot) = beamWidth; - decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; runtime::kernels::invokeFill(*endIdTensorPtr, endId, decoderStream); @@ -320,15 +318,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder beamHypotheses.init(manager, endId); } - // Speculative execution - if (numDecodingEngineTokens > 1 || decoderState.getSpeculativeDecodingMode().isDraftTokensExternal()) - { - TLLM_CHECK(beamWidth == 1); - newRequestSpeculativeDecoding(batchSlot, request, samplingConfig, modelConfig, - decoderState.getJointDecodingInput(), decoderState.getJointDecodingOutput(), runtimeStream, decoderStream, - decoderState.getSpeculativeDecodingMode(), decoderState.getMaxDecodingEngineTokens()); - } - // fill outputIds with endIds TensorPtr const outputIds = ITensor::slice(dJointOutput.ids, batchSlot, 1); auto outputIdsTileView = ITensor::view(outputIds, ITensor::makeShape({beamWidth, maxSequenceLength})); @@ -608,6 +597,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon auto decoderRequest = decoder_batch::Request{inputView, promptLen, llmReq->mMaxNewTokens, llmReq->mEndId}; llmReq->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; + if (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) { if (llmReq->hasDraftTokens()) @@ -653,6 +643,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon decoderRequest.eagleConfig = llmReq->getEagleConfig() ? llmReq->getEagleConfig() : decodingConfig.getEagleConfig(); } + if (llmReq->getEmbeddingBias().has_value()) { decoderRequest.embeddingBias = getEmbeddingBias(logitsType, llmReq->getEmbeddingBias().value()); @@ -671,8 +662,25 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon } TLLM_CHECK(llmReq->mSeqSlot.has_value()); - newRequest(llmReq->mSeqSlot.value(), decoderRequest, llmReq->mSamplingConfig, modelConfig, decoderState, - runtimeStream, decoderStream, maxSequenceLength); + auto const batchSlot = llmReq->mSeqSlot.value(); + auto const& samplingConfig = llmReq->mSamplingConfig; + + newRequest( + batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState, decoderStream, maxSequenceLength); + + auto const numDecodingEngineTokens = decoderRequest.generatedTokensPerEngineStep; + decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); + + // Speculative execution + if (numDecodingEngineTokens > 1 || decoderState.getSpeculativeDecodingMode().isDraftTokensExternal()) + { + auto const beamWidth = samplingConfig.beamWidth; + TLLM_CHECK(beamWidth == 1); + + newRequestSpeculativeDecoding(batchSlot, decoderRequest, samplingConfig, modelConfig, + decoderState.getJointDecodingInput(), decoderState.getJointDecodingOutput(), runtimeStream, + decoderStream, decoderState.getSpeculativeDecodingMode(), decoderState.getMaxDecodingEngineTokens()); + } decoderRequests.push_back(decoderRequest); From f84697750a5ca97069412d192346cb7f102fb0f9 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 11:52:19 +0000 Subject: [PATCH 06/17] refactor: Consolidate speculative decoding logic in createDecoderRequests - Moved the speculative decoding logic for Medusa, Lookahead, and Eagle modes into a unified conditional block for improved clarity and maintainability. - Ensured that the handling of decoder requests is streamlined and easier to follow. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 47 ++++++++++--------- 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 7964a279d3c7..963ad8d80e35 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -621,28 +621,6 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon { decoderRequest.generatedTokensPerEngineStep = modelConfig.getMaxDecodingTokens(); } - if (modelConfig.getSpeculativeDecodingMode().isMedusa()) - { - TLLM_CHECK(medusaBuffers); - llmReq->mSamplingConfig.topKMedusaHeads = {medusaBuffers->mTopKs}; - // FIXME: we must set medusa paths and tree ids not from seq slot, but from llmRequest? - // When multiple microbatches buffers are used, runtime buffers can not be addressed with seqSlot. - decoderRequest.medusaPaths = ITensor::slice(medusaBuffers->medusaPathsDevice, 0, 1); - decoderRequest.medusaTreeIds = ITensor::slice(medusaBuffers->medusaTreeIdsDevice, 0, 1); - } - else if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - lookaheadPrompt.emplace_back(ITensor::slice(decoderRequest.ids, 0, decoderRequest.inputLen)); - - auto const& lookaheadRuntimeConfig - = llmReq->getLookaheadConfig().value_or(decodingConfig.getLookaheadDecodingConfig().value()); - lookaheadAlgoConfigs.emplace_back(lookaheadRuntimeConfig); - } - else if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - decoderRequest.eagleConfig - = llmReq->getEagleConfig() ? llmReq->getEagleConfig() : decodingConfig.getEagleConfig(); - } if (llmReq->getEmbeddingBias().has_value()) { @@ -672,11 +650,34 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); // Speculative execution - if (numDecodingEngineTokens > 1 || decoderState.getSpeculativeDecodingMode().isDraftTokensExternal()) + if (!decoderState.getSpeculativeDecodingMode().isNone()) { auto const beamWidth = samplingConfig.beamWidth; TLLM_CHECK(beamWidth == 1); + if (modelConfig.getSpeculativeDecodingMode().isMedusa()) + { + TLLM_CHECK(medusaBuffers); + llmReq->mSamplingConfig.topKMedusaHeads = {medusaBuffers->mTopKs}; + // FIXME: we must set medusa paths and tree ids not from seq slot, but from llmRequest? + // When multiple microbatches buffers are used, runtime buffers can not be addressed with seqSlot. + decoderRequest.medusaPaths = ITensor::slice(medusaBuffers->medusaPathsDevice, 0, 1); + decoderRequest.medusaTreeIds = ITensor::slice(medusaBuffers->medusaTreeIdsDevice, 0, 1); + } + else if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) + { + lookaheadPrompt.emplace_back(ITensor::slice(decoderRequest.ids, 0, decoderRequest.inputLen)); + + auto const& lookaheadRuntimeConfig + = llmReq->getLookaheadConfig().value_or(decodingConfig.getLookaheadDecodingConfig().value()); + lookaheadAlgoConfigs.emplace_back(lookaheadRuntimeConfig); + } + else if (modelConfig.getSpeculativeDecodingMode().isEagle()) + { + decoderRequest.eagleConfig + = llmReq->getEagleConfig() ? llmReq->getEagleConfig() : decodingConfig.getEagleConfig(); + } + newRequestSpeculativeDecoding(batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState.getJointDecodingInput(), decoderState.getJointDecodingOutput(), runtimeStream, decoderStream, decoderState.getSpeculativeDecodingMode(), decoderState.getMaxDecodingEngineTokens()); From 60f266feed11ead8e874c95ce639ae37156855aa Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 12:01:51 +0000 Subject: [PATCH 07/17] refactor: Enhance embedding bias initialization in CreateNewDecoderRequests - Updated the initializeEmbeddingBias function to accept LlmRequest and logitsType, improving clarity and flexibility. - Removed direct embedding bias handling from newRequest, streamlining the request creation process. - Ensured that embedding bias is now consistently initialized using the updated function. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/request.h | 1 - .../createNewDecoderRequests.cpp | 23 +++++++++++-------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/request.h b/cpp/include/tensorrt_llm/runtime/request.h index 1861ea843174..e9c985156200 100644 --- a/cpp/include/tensorrt_llm/runtime/request.h +++ b/cpp/include/tensorrt_llm/runtime/request.h @@ -48,7 +48,6 @@ class Request std::optional maxNewTokens; // maximum number of tokens to generate for this request std::optional endId; // end token id SizeType32 generatedTokensPerEngineStep{1}; // - TensorPtr embeddingBias; // [vocabSizePadded], on gpu TensorPtr badWordsList; // [2, badWordsLength] on gpu TensorPtr stopWordsList; // [2, stopWordsLength] on gpu diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 963ad8d80e35..7e96fa8491b5 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -170,12 +170,15 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru namespace { -void initializeEmbeddingBias(DecodingInput& dJointInput, TensorPtr const& embeddingBiasTensor, SizeType32 batchSlot, +void initializeEmbeddingBias(DecodingInput& dJointInput, SizeType32 batchSlot, + std::optional const& embeddingBias, nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, BufferManager const& manager) { TensorPtr const embeddingBiasSlice = ITensor::slice(constPointerCast(dJointInput.embeddingBias), batchSlot, 1); - if (embeddingBiasTensor) + if (embeddingBias.has_value()) { + auto embeddingBiasTensor = getEmbeddingBias(logitsType, embeddingBias.value()); + TLLM_CHECK(embeddingBiasTensor->getShape().nbDims == 2); TLLM_CHECK(embeddingBiasTensor->getShape().d[0] == 1); TLLM_CHECK_WITH_INFO(embeddingBiasTensor->getShape().d[1] == modelConfig.getVocabSize(), @@ -252,8 +255,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; runtime::kernels::invokeFill(*endIdTensorPtr, endId, decoderStream); - initializeEmbeddingBias(dJointInput, request.embeddingBias, batchSlot, modelConfig, manager); - setupWords(dJointInput.stopWordsLists, request.stopWordsList, dJointInput.stopWordsPtrs, dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot); @@ -588,6 +589,11 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon SizeType32 inputOffset{0}; for (auto const& llmReq : finishedContextRequests) { + auto& dJointInput = decoderState.getJointDecodingInput(); + + TLLM_CHECK(llmReq->mSeqSlot.has_value()); + auto const batchSlot = llmReq->mSeqSlot.value(); + auto const promptLen = llmReq->getPromptLen(); auto const& reqTokens = llmReq->getTokens(0); TLLM_CHECK(reqTokens.size() == static_cast(promptLen)); @@ -622,10 +628,9 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon decoderRequest.generatedTokensPerEngineStep = modelConfig.getMaxDecodingTokens(); } - if (llmReq->getEmbeddingBias().has_value()) - { - decoderRequest.embeddingBias = getEmbeddingBias(logitsType, llmReq->getEmbeddingBias().value()); - } + initializeEmbeddingBias( + dJointInput, batchSlot, llmReq->getEmbeddingBias(), logitsType, modelConfig, bufferManager); + if (llmReq->getBadWordsList().has_value()) { // Move to GPU and remove leading bs1 dimension since this is what decoderRequest expects @@ -639,8 +644,6 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon decoderRequest.stopWordsList->squeeze(0); } - TLLM_CHECK(llmReq->mSeqSlot.has_value()); - auto const batchSlot = llmReq->mSeqSlot.value(); auto const& samplingConfig = llmReq->mSamplingConfig; newRequest( From c87332ceef10911c8a2a86e58c354fdca2d23faf Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 12:22:25 +0000 Subject: [PATCH 08/17] refactor: Update stopWords and badWords handling in CreateNewDecoderRequests - Removed badWordsList and stopWordsList from the request structure, streamlining the request parameters. - Enhanced setupWords function to handle optional TensorPtr for requestWordsList, improving flexibility. - Updated calls to setupWords in newRequest to utilize the new structure, ensuring proper memory management and GPU handling. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/request.h | 2 - .../createNewDecoderRequests.cpp | 40 ++++++++----------- 2 files changed, 16 insertions(+), 26 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/request.h b/cpp/include/tensorrt_llm/runtime/request.h index e9c985156200..9a5e0501f094 100644 --- a/cpp/include/tensorrt_llm/runtime/request.h +++ b/cpp/include/tensorrt_llm/runtime/request.h @@ -48,8 +48,6 @@ class Request std::optional maxNewTokens; // maximum number of tokens to generate for this request std::optional endId; // end token id SizeType32 generatedTokensPerEngineStep{1}; // - TensorPtr badWordsList; // [2, badWordsLength] on gpu - TensorPtr stopWordsList; // [2, stopWordsLength] on gpu //! Optional parameters for speculative decoding BufferPtr draftTokens; // [generatedTokensPerEngineStep - 1] on gpu diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 7e96fa8491b5..7d890a248905 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -192,21 +192,26 @@ void initializeEmbeddingBias(DecodingInput& dJointInput, SizeType32 batchSlot, } } -void setupWords(std::vector& jointWordsLists, TensorPtr const& requestWordsList, - SharedConstPtr& jointWordsPtrs, SharedConstPtr& jointWordsLens, SizeType32& jointMaxWordsLen, SizeType32 batchSlot) +void setupWords(std::vector& jointWordsLists, + std::optional const& requestWordsList, SharedConstPtr& jointWordsPtrs, SharedConstPtr& jointWordsLens, + SizeType32& jointMaxWordsLen, SizeType32 batchSlot, BufferManager const& manager) { - if (requestWordsList) + if (requestWordsList.has_value()) { - auto const wordsLen = requestWordsList->getShape().d[1]; + // Move to GPU and remove leading bs1 dimension since this is what decoderRequest expects + TensorPtr wordsList = manager.copyFrom(*requestWordsList.value(), MemoryType::kGPU); + wordsList->squeeze(0); + + auto const wordsLen = wordsList->getShape().d[1]; BufferRange(*constPointerCast(jointWordsPtrs))[batchSlot] - = runtime::bufferCast(*requestWordsList); + = runtime::bufferCast(*wordsList); runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = wordsLen; // FIXME: this is monotonically growing size jointMaxWordsLen = std::max(static_cast(wordsLen), jointMaxWordsLen); // NOTE: jointWordsList is not used in gptDecoder, but required to keep WordsList's // memory allocated - jointWordsLists[batchSlot] = requestWordsList; + jointWordsLists[batchSlot] = wordsList; } else { @@ -255,12 +260,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; runtime::kernels::invokeFill(*endIdTensorPtr, endId, decoderStream); - setupWords(dJointInput.stopWordsLists, request.stopWordsList, dJointInput.stopWordsPtrs, dJointInput.stopWordsLens, - dJointInput.maxStopWordsLen, batchSlot); - - setupWords(dJointInput.badWordsLists, request.badWordsList, dJointInput.badWordsPtrs, dJointInput.badWordsLens, - dJointInput.maxBadWordsLen, batchSlot); - TensorPtr const sequenceLimitLength{ ITensor::slice(constPointerCast(dJointInput.sequenceLimitLength), batchSlot, 1)}; runtime::kernels::invokeFill(*sequenceLimitLength, inputLength + maxNewTokens, decoderStream); @@ -631,18 +630,11 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon initializeEmbeddingBias( dJointInput, batchSlot, llmReq->getEmbeddingBias(), logitsType, modelConfig, bufferManager); - if (llmReq->getBadWordsList().has_value()) - { - // Move to GPU and remove leading bs1 dimension since this is what decoderRequest expects - decoderRequest.badWordsList = bufferManager.copyFrom(*llmReq->getBadWordsList().value(), MemoryType::kGPU); - decoderRequest.badWordsList->squeeze(0); - } - if (llmReq->getStopWordsList().has_value()) - { - decoderRequest.stopWordsList - = bufferManager.copyFrom(*llmReq->getStopWordsList().value(), MemoryType::kGPU); - decoderRequest.stopWordsList->squeeze(0); - } + setupWords(dJointInput.badWordsLists, llmReq->getBadWordsList(), dJointInput.badWordsPtrs, + dJointInput.badWordsLens, dJointInput.maxBadWordsLen, batchSlot, bufferManager); + + setupWords(dJointInput.stopWordsLists, llmReq->getStopWordsList(), dJointInput.stopWordsPtrs, + dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, bufferManager); auto const& samplingConfig = llmReq->mSamplingConfig; From 7a3e5e6600b11571d438a42ea0b0024eb3c75194 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:03:26 +0000 Subject: [PATCH 09/17] refactor: Introduce initializeRequestIds and initializeBeamSearch functions in CreateNewDecoderRequests - Added initializeRequestIds function to handle the initialization of request IDs and output IDs, improving code organization. - Introduced initializeBeamSearch function to encapsulate beam search initialization logic, enhancing clarity and maintainability. - Updated newRequest to utilize the new functions, streamlining the request handling process. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 74 +++++++++++-------- 1 file changed, 44 insertions(+), 30 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 7d890a248905..cb01e144ae09 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -170,6 +170,46 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru namespace { +void initializeRequestIds(DecodingInput& dJointInput, DecodingOutput& dJointOutput, SizeType32 batchSlot, + SharedConstPtr const& requestIds, SizeType32 endId, SizeType32 beamWidth, SizeType32 maxSequenceLength, + BufferManager const& manager) +{ + TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; + runtime::kernels::invokeFill(*endIdTensorPtr, endId, manager.getStream()); + + // fill outputIds with endIds + TensorPtr const outputIds = ITensor::slice(dJointOutput.ids, batchSlot, 1); + auto outputIdsTileView = ITensor::view(outputIds, ITensor::makeShape({beamWidth, maxSequenceLength})); + runtime::kernels::invokeFill(*outputIdsTileView, endId, manager.getStream()); + + // copy the request ids into outputIds + auto const requestIdsShape = requestIds->getShape(); + auto outputIdsView = ITensor::view(outputIds, requestIdsShape); + manager.copy(*requestIds, *outputIdsView); +} + +void initializeBeamSearch(DecodingInput& dJointInput, DecodingOutput& dJointOutput, SizeType32 batchSlot, + SizeType32 endId, SizeType32 beamWidth, SizeType32 maxSequenceLength, BufferManager const& manager) +{ + TensorPtr const cumLogProbs = ITensor::slice(dJointOutput.cumLogProbs, batchSlot, 1); + runtime::kernels::invokeFill( + *IBuffer::slice(cumLogProbs, 1, beamWidth - 1), DecodingOutput::kNegativeInfinity, manager.getStream()); + + auto parentIds = ITensor::slice(dJointOutput.parentIds, batchSlot, 1); + auto const outputIdsShape = ITensor::makeShape({1, beamWidth, maxSequenceLength}); + parentIds->reshape(outputIdsShape); + manager.setZero(*parentIds); + + auto cacheIndirectionInput = ITensor::slice(dJointInput.cacheIndirection, batchSlot, 1); + manager.setZero(*cacheIndirectionInput); + + auto cacheIndirectionOutput = ITensor::slice(dJointOutput.cacheIndirection, batchSlot, 1); + manager.setZero(*cacheIndirectionOutput); + + auto beamHypotheses = dJointOutput.beamHypotheses.slice(batchSlot, 1); + beamHypotheses.init(manager, endId); +} + void initializeEmbeddingBias(DecodingInput& dJointInput, SizeType32 batchSlot, std::optional const& embeddingBias, nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, BufferManager const& manager) @@ -257,9 +297,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder dJointInput.beamWidths.at(batchSlot) = beamWidth; - TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; - runtime::kernels::invokeFill(*endIdTensorPtr, endId, decoderStream); - TensorPtr const sequenceLimitLength{ ITensor::slice(constPointerCast(dJointInput.sequenceLimitLength), batchSlot, 1)}; runtime::kernels::invokeFill(*sequenceLimitLength, inputLength + maxNewTokens, decoderStream); @@ -269,7 +306,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder // output auto& dJointOutput = decoderState.getJointDecodingOutput(); - auto const outputIdsShape = ITensor::makeShape({1, beamWidth, maxSequenceLength}); auto finishedSum = ITensor::slice(dJointOutput.finishedSum, batchSlot, 1); manager.setZero(*finishedSum); @@ -298,36 +334,14 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder manager.setZero(*logProbs); } + initializeRequestIds( + dJointInput, dJointOutput, batchSlot, requestIds, endId, beamWidth, maxSequenceLength, manager); + if (beamWidth > 1) { - TensorPtr const cumLogProbs = ITensor::slice(dJointOutput.cumLogProbs, batchSlot, 1); - runtime::kernels::invokeFill( - *IBuffer::slice(cumLogProbs, 1, beamWidth - 1), DecodingOutput::kNegativeInfinity, decoderStream); - - auto parentIds = ITensor::slice(dJointOutput.parentIds, batchSlot, 1); - parentIds->reshape(outputIdsShape); - manager.setZero(*parentIds); - - auto cacheIndirectionInput = ITensor::slice(dJointInput.cacheIndirection, batchSlot, 1); - manager.setZero(*cacheIndirectionInput); - - auto cacheIndirectionOutput = ITensor::slice(dJointOutput.cacheIndirection, batchSlot, 1); - manager.setZero(*cacheIndirectionOutput); - - auto beamHypotheses = dJointOutput.beamHypotheses.slice(batchSlot, 1); - beamHypotheses.init(manager, endId); + initializeBeamSearch(dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, manager); } - // fill outputIds with endIds - TensorPtr const outputIds = ITensor::slice(dJointOutput.ids, batchSlot, 1); - auto outputIdsTileView = ITensor::view(outputIds, ITensor::makeShape({beamWidth, maxSequenceLength})); - runtime::kernels::invokeFill(*outputIdsTileView, endId, decoderStream); - - // copy the request ids into outputIds - auto const requestIdsShape = requestIds->getShape(); - auto outputIdsView = ITensor::view(outputIds, requestIdsShape); - manager.copy(*requestIds, *outputIdsView); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } From c8c94a88fcf7c174e8d419bb49d5585fc0b1d5f1 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 12:50:14 +0000 Subject: [PATCH 10/17] refactor: Remove ids and endId from decoder_batch::Request - Modified the Request constructor to remove the ids and endId parameters, streamlining the initialization process. - Updated the newRequest function in CreateNewDecoderRequests to reflect changes in the Request structure, ensuring proper handling of input length and token generation. - Removed unnecessary checks related to request IDs and endId, enhancing code clarity and maintainability. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/request.h | 9 +--- .../createNewDecoderRequests.cpp | 43 ++++++++++--------- 2 files changed, 25 insertions(+), 27 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/request.h b/cpp/include/tensorrt_llm/runtime/request.h index 9a5e0501f094..983d19ad8c98 100644 --- a/cpp/include/tensorrt_llm/runtime/request.h +++ b/cpp/include/tensorrt_llm/runtime/request.h @@ -31,22 +31,17 @@ class Request using TensorPtr = ITensor::SharedPtr; using BufferPtr = IBuffer::SharedPtr; - explicit Request(TensorConstPtr ids, SizeType32 inputLen, std::optional maxNewTokens = std::nullopt, - std::optional endId = std::nullopt) - : ids{std::move(ids)} - , inputLen(inputLen) + explicit Request(SizeType32 inputLen, std::optional maxNewTokens = std::nullopt) + : inputLen(inputLen) , maxNewTokens{maxNewTokens} - , endId{endId} { } //! Mandatory parameters - TensorConstPtr ids; // The input sequence of token ids, [inputSeqLen], on gpu SizeType32 inputLen; // Input length without draft tokens, increasing with generation steps // optional parameters std::optional maxNewTokens; // maximum number of tokens to generate for this request - std::optional endId; // end token id SizeType32 generatedTokensPerEngineStep{1}; // //! Optional parameters for speculative decoding diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index cb01e144ae09..d278b41fdb30 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -278,7 +278,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder TLLM_CHECK_WITH_INFO(beamWidth <= maxBeamWidth, tc::fmtstr("Beam width (%d) must be smaller than maxBeamWidth (%d) passed to decoder setup function.", beamWidth, maxBeamWidth)); - auto const& requestIds = request.ids; auto const inputLength = request.inputLen; auto const numDecodingEngineTokens = request.generatedTokensPerEngineStep; auto const numDecodingDraftEngineTokens = numDecodingEngineTokens - 1; @@ -289,8 +288,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder tc::fmtstr( "Input length (%d) + max new tokens (%d) + draft tokens (%d) must be less than max sequence length (%d).", inputLength, maxNewTokens, numDecodingDraftEngineTokens, maxSequenceLength)); - TLLM_CHECK(requestIds->getDataType() == TRTDataType::value); - auto const endId = request.endId.value_or(-1); // input auto& dJointInput = decoderState.getJointDecodingInput(); @@ -334,14 +331,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder manager.setZero(*logProbs); } - initializeRequestIds( - dJointInput, dJointOutput, batchSlot, requestIds, endId, beamWidth, maxSequenceLength, manager); - - if (beamWidth > 1) - { - initializeBeamSearch(dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, manager); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } @@ -607,13 +596,12 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon TLLM_CHECK(llmReq->mSeqSlot.has_value()); auto const batchSlot = llmReq->mSeqSlot.value(); + auto const& samplingConfig = llmReq->mSamplingConfig; + auto const beamWidth = samplingConfig.beamWidth; + auto const promptLen = llmReq->getPromptLen(); - auto const& reqTokens = llmReq->getTokens(0); - TLLM_CHECK(reqTokens.size() == static_cast(promptLen)); - TensorPtr inputView = ITensor::slice(inputIds, inputOffset, promptLen); - bufferManager.copy(reqTokens.data(), *inputView); - auto decoderRequest = decoder_batch::Request{inputView, promptLen, llmReq->mMaxNewTokens, llmReq->mEndId}; + auto decoderRequest = decoder_batch::Request{promptLen, llmReq->mMaxNewTokens}; llmReq->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; @@ -650,18 +638,33 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon setupWords(dJointInput.stopWordsLists, llmReq->getStopWordsList(), dJointInput.stopWordsPtrs, dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, bufferManager); - auto const& samplingConfig = llmReq->mSamplingConfig; - newRequest( batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState, decoderStream, maxSequenceLength); + auto& dJointOutput = decoderState.getJointDecodingOutput(); + + auto const& reqTokens = llmReq->getTokens(0); + TLLM_CHECK(reqTokens.size() == static_cast(promptLen)); + TensorPtr requestIds = ITensor::slice(inputIds, inputOffset, promptLen); + // Copy to pinned host memory (don't care about stream of bufferManager) + bufferManager.copy(reqTokens.data(), *requestIds); + auto const endId = llmReq->mEndId.value_or(-1); + + initializeRequestIds( + dJointInput, dJointOutput, batchSlot, requestIds, endId, beamWidth, maxSequenceLength, bufferManager); + + if (beamWidth > 1) + { + initializeBeamSearch( + dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, bufferManager); + } + auto const numDecodingEngineTokens = decoderRequest.generatedTokensPerEngineStep; decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); // Speculative execution if (!decoderState.getSpeculativeDecodingMode().isNone()) { - auto const beamWidth = samplingConfig.beamWidth; TLLM_CHECK(beamWidth == 1); if (modelConfig.getSpeculativeDecodingMode().isMedusa()) @@ -675,7 +678,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon } else if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) { - lookaheadPrompt.emplace_back(ITensor::slice(decoderRequest.ids, 0, decoderRequest.inputLen)); + lookaheadPrompt.emplace_back(requestIds); auto const& lookaheadRuntimeConfig = llmReq->getLookaheadConfig().value_or(decodingConfig.getLookaheadDecodingConfig().value()); From 2000700bdb3dc1e715f96ca50ece9fb80eef3914 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:23:32 +0000 Subject: [PATCH 11/17] refactor: Introduce initializeLogProbs function in CreateNewDecoderRequests - Added initializeLogProbs function to encapsulate the initialization of cumulative and output log probabilities, improving code organization and clarity. - Removed redundant log probability initialization logic from newRequest, streamlining the request handling process. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 34 ++++++++++++------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index d278b41fdb30..6f1020f58100 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -259,6 +259,25 @@ void setupWords(std::vector& jointWordsLists, } }; +void initializeLogProbs(DecodingOutput& dJointOutput, SizeType32 batchSlot, SamplingConfig const& samplingConfig, + BufferManager const& manager) +{ + auto const beamWidth = samplingConfig.beamWidth; + + // cumLogProb is mandatory for beamWidth > 1 + if ((samplingConfig.cumLogProbs.has_value() && samplingConfig.cumLogProbs->at(0)) || beamWidth > 1) + { + auto cumLogProbs = ITensor::slice(dJointOutput.cumLogProbs, batchSlot, 1); + manager.setZero(*cumLogProbs); + } + + if (samplingConfig.outputLogProbs.has_value() && samplingConfig.outputLogProbs->at(0)) + { + auto logProbs = ITensor::slice(dJointOutput.logProbs, batchSlot, 1); + manager.setZero(*logProbs); + } +} + } // namespace void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, @@ -318,19 +337,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder TensorPtr const finishedStepsSlice = ITensor::slice(decoderState.getFinishReasons(), batchSlot, 1); manager.setZero(*finishedStepsSlice); - // cumLogProb is mandatory for beamWidth > 1 - if ((samplingConfig.cumLogProbs.has_value() && samplingConfig.cumLogProbs->at(0)) || beamWidth > 1) - { - auto cumLogProbs = ITensor::slice(dJointOutput.cumLogProbs, batchSlot, 1); - manager.setZero(*cumLogProbs); - } - - if (samplingConfig.outputLogProbs.has_value() && samplingConfig.outputLogProbs->at(0)) - { - auto logProbs = ITensor::slice(dJointOutput.logProbs, batchSlot, 1); - manager.setZero(*logProbs); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } @@ -643,6 +649,8 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon auto& dJointOutput = decoderState.getJointDecodingOutput(); + initializeLogProbs(dJointOutput, batchSlot, samplingConfig, bufferManager); + auto const& reqTokens = llmReq->getTokens(0); TLLM_CHECK(reqTokens.size() == static_cast(promptLen)); TensorPtr requestIds = ITensor::slice(inputIds, inputOffset, promptLen); From 163a93c102c277a898a2dc4bb3966eff9c3462e7 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:24:55 +0000 Subject: [PATCH 12/17] refactor: Use decoder stream for request initialization - Fix: Use decoder stream instead of runtime stream for request initialization. - Updated functions in CreateNewDecoderRequests to utilize BufferManager instead of CudaStream for stream management, enhancing code clarity and consistency. - Streamlined the initialization of request IDs, beam search, and embedding bias by passing BufferManager, improving memory management and GPU handling. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 6f1020f58100..cfc66b03e512 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -571,7 +571,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon runtime::CudaStream const& runtimeStream, runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const { - auto const bufferManager = BufferManager{std::make_shared(runtimeStream.get())}; + auto const decoderBufferManager = BufferManager{std::make_shared(decoderStream.get())}; unsigned decoderInputSize{0}; for (auto const& llmReq : finishedContextRequests) @@ -616,12 +616,13 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon if (llmReq->hasDraftTokens()) { auto const& draftTokens = llmReq->getDraftTokens(); - decoderRequest.draftTokens = bufferManager.copyFrom(*draftTokens, MemoryType::kPINNEDPOOL); + // Copy to pinned host memory (don't care about stream of bufferManager) + decoderRequest.draftTokens = decoderBufferManager.copyFrom(*draftTokens, MemoryType::kPINNEDPOOL); auto const& draftLogits = llmReq->getDraftLogits(); if (draftLogits.has_value()) { decoderRequest.draftLogits - = retrieveDraftLogits(modelConfig, worldConfig, draftLogits.value(), bufferManager); + = retrieveDraftLogits(modelConfig, worldConfig, draftLogits.value(), decoderBufferManager); } decoderRequest.generatedTokensPerEngineStep = draftTokens->size() + 1; } @@ -636,35 +637,35 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon } initializeEmbeddingBias( - dJointInput, batchSlot, llmReq->getEmbeddingBias(), logitsType, modelConfig, bufferManager); + dJointInput, batchSlot, llmReq->getEmbeddingBias(), logitsType, modelConfig, decoderBufferManager); setupWords(dJointInput.badWordsLists, llmReq->getBadWordsList(), dJointInput.badWordsPtrs, - dJointInput.badWordsLens, dJointInput.maxBadWordsLen, batchSlot, bufferManager); + dJointInput.badWordsLens, dJointInput.maxBadWordsLen, batchSlot, decoderBufferManager); setupWords(dJointInput.stopWordsLists, llmReq->getStopWordsList(), dJointInput.stopWordsPtrs, - dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, bufferManager); + dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, decoderBufferManager); newRequest( batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState, decoderStream, maxSequenceLength); auto& dJointOutput = decoderState.getJointDecodingOutput(); - initializeLogProbs(dJointOutput, batchSlot, samplingConfig, bufferManager); + initializeLogProbs(dJointOutput, batchSlot, samplingConfig, decoderBufferManager); auto const& reqTokens = llmReq->getTokens(0); TLLM_CHECK(reqTokens.size() == static_cast(promptLen)); TensorPtr requestIds = ITensor::slice(inputIds, inputOffset, promptLen); // Copy to pinned host memory (don't care about stream of bufferManager) - bufferManager.copy(reqTokens.data(), *requestIds); + decoderBufferManager.copy(reqTokens.data(), *requestIds); auto const endId = llmReq->mEndId.value_or(-1); - initializeRequestIds( - dJointInput, dJointOutput, batchSlot, requestIds, endId, beamWidth, maxSequenceLength, bufferManager); + initializeRequestIds(dJointInput, dJointOutput, batchSlot, requestIds, endId, beamWidth, maxSequenceLength, + decoderBufferManager); if (beamWidth > 1) { initializeBeamSearch( - dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, bufferManager); + dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, decoderBufferManager); } auto const numDecodingEngineTokens = decoderRequest.generatedTokensPerEngineStep; From efe0aab17ac5298cfcf56de10c7a912a4e5c315e Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:59:49 +0000 Subject: [PATCH 13/17] refactor: Move beam width check to CreateNewDecoderRequests - Move beam width check to CreateNewDecoderRequests from newRequest. - Set beam width in CreateNewDecoderRequests instead of newRequest. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.cpp | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index cfc66b03e512..8880242eb3a4 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -292,11 +292,7 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder auto const batchSize = decoderState.getMaxNumSequences(); TLLM_CHECK(0 <= batchSize && batchSlot < batchSize); - auto const maxBeamWidth = decoderState.getMaxBeamWidth(); - auto const beamWidth = samplingConfig.beamWidth; - TLLM_CHECK_WITH_INFO(beamWidth <= maxBeamWidth, - tc::fmtstr("Beam width (%d) must be smaller than maxBeamWidth (%d) passed to decoder setup function.", - beamWidth, maxBeamWidth)); + auto const inputLength = request.inputLen; auto const numDecodingEngineTokens = request.generatedTokensPerEngineStep; auto const numDecodingDraftEngineTokens = numDecodingEngineTokens - 1; @@ -311,8 +307,6 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder // input auto& dJointInput = decoderState.getJointDecodingInput(); - dJointInput.beamWidths.at(batchSlot) = beamWidth; - TensorPtr const sequenceLimitLength{ ITensor::slice(constPointerCast(dJointInput.sequenceLimitLength), batchSlot, 1)}; runtime::kernels::invokeFill(*sequenceLimitLength, inputLength + maxNewTokens, decoderStream); @@ -597,20 +591,26 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon SizeType32 inputOffset{0}; for (auto const& llmReq : finishedContextRequests) { + llmReq->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; + auto& dJointInput = decoderState.getJointDecodingInput(); TLLM_CHECK(llmReq->mSeqSlot.has_value()); auto const batchSlot = llmReq->mSeqSlot.value(); auto const& samplingConfig = llmReq->mSamplingConfig; + auto const beamWidth = samplingConfig.beamWidth; + auto const maxBeamWidth = decoderState.getMaxBeamWidth(); + TLLM_CHECK_WITH_INFO(beamWidth <= maxBeamWidth, + tc::fmtstr("Beam width (%d) must be smaller than maxBeamWidth (%d) passed to decoder setup function.", + beamWidth, maxBeamWidth)); + dJointInput.beamWidths.at(batchSlot) = beamWidth; auto const promptLen = llmReq->getPromptLen(); auto decoderRequest = decoder_batch::Request{promptLen, llmReq->mMaxNewTokens}; - llmReq->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; - if (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) { if (llmReq->hasDraftTokens()) From 4b2df1ec593c0c079e495ad9660faa237c0718f0 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 14:02:05 +0000 Subject: [PATCH 14/17] refactor: Introduce initializeInputLengths function in CreateNewDecoderRequests - Added initializeInputLengths function to encapsulate the logic for setting input lengths and validating constraints, improving code organization and clarity. - Streamlined the newRequest function by replacing direct input length handling with a call to initializeInputLengths, enhancing maintainability. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../createNewDecoderRequests.cpp | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 8880242eb3a4..5bd20aade108 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -170,6 +170,26 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru namespace { +void initializeInputLengths(DecodingInput& dJointInput, SizeType32 batchSlot, SizeType32 inputLength, + std::optional maxNewTokensOpt, SizeType32 numDecodingEngineTokens, SizeType32 maxSequenceLength, + BufferManager const& manager) +{ + auto const numDecodingDraftEngineTokens = numDecodingEngineTokens - 1; + auto const maxNewTokens = maxNewTokensOpt.value_or(maxSequenceLength - inputLength - numDecodingDraftEngineTokens); + + TLLM_CHECK_WITH_INFO(inputLength + maxNewTokens + numDecodingDraftEngineTokens <= maxSequenceLength, + tc::fmtstr( + "Input length (%d) + max new tokens (%d) + draft tokens (%d) must be less than max sequence length (%d).", + inputLength, maxNewTokens, numDecodingDraftEngineTokens, maxSequenceLength)); + + TensorPtr const sequenceLimitLength{ + ITensor::slice(constPointerCast(dJointInput.sequenceLimitLength), batchSlot, 1)}; + runtime::kernels::invokeFill(*sequenceLimitLength, inputLength + maxNewTokens, manager.getStream()); + + TensorPtr const inputLengths{ITensor::slice(constPointerCast(dJointInput.lengths), batchSlot, 1)}; + runtime::kernels::invokeFill(*inputLengths, inputLength, manager.getStream()); +} + void initializeRequestIds(DecodingInput& dJointInput, DecodingOutput& dJointOutput, SizeType32 batchSlot, SharedConstPtr const& requestIds, SizeType32 endId, SizeType32 beamWidth, SizeType32 maxSequenceLength, BufferManager const& manager) @@ -295,24 +315,12 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder auto const inputLength = request.inputLen; auto const numDecodingEngineTokens = request.generatedTokensPerEngineStep; - auto const numDecodingDraftEngineTokens = numDecodingEngineTokens - 1; - auto const maxNewTokens - = request.maxNewTokens.value_or(maxSequenceLength - inputLength - numDecodingDraftEngineTokens); - - TLLM_CHECK_WITH_INFO(inputLength + maxNewTokens + numDecodingDraftEngineTokens <= maxSequenceLength, - tc::fmtstr( - "Input length (%d) + max new tokens (%d) + draft tokens (%d) must be less than max sequence length (%d).", - inputLength, maxNewTokens, numDecodingDraftEngineTokens, maxSequenceLength)); // input auto& dJointInput = decoderState.getJointDecodingInput(); - TensorPtr const sequenceLimitLength{ - ITensor::slice(constPointerCast(dJointInput.sequenceLimitLength), batchSlot, 1)}; - runtime::kernels::invokeFill(*sequenceLimitLength, inputLength + maxNewTokens, decoderStream); - - TensorPtr const inputLengths{ITensor::slice(constPointerCast(dJointInput.lengths), batchSlot, 1)}; - runtime::kernels::invokeFill(*inputLengths, inputLength, decoderStream); + initializeInputLengths(dJointInput, batchSlot, inputLength, request.maxNewTokens, + request.generatedTokensPerEngineStep, maxSequenceLength, manager); // output auto& dJointOutput = decoderState.getJointDecodingOutput(); From a4f9c9705408b5393055ff38eba7c27ec78a3b1e Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 14:19:33 +0000 Subject: [PATCH 15/17] refactor: Initialize request lengths in CreateNewDecoderRequests - Moved the initialization of input lengths from newRequest to the createDecoderRequests function. - Removed the now unnecessary maxNewTokens parameter from the Request constructor, simplifying its initialization. - Removed the now unnecessary maxSequenceLength parameter from the newRequest function. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.h | 2 +- cpp/include/tensorrt_llm/runtime/request.h | 4 +-- .../createNewDecoderRequests.cpp | 31 ++++++------------- 3 files changed, 12 insertions(+), 25 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index c35ba95d5a0a..e129cf6e9ab6 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -91,7 +91,7 @@ class CreateNewDecoderRequests : Algorithm //! @brief Initialize the decoder at `batchSlot` with a new `request`. static void newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, - runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream, SizeType32 maxSequenceLength); + runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream); //! @brief Setups decoder internal tensors for new speculative decoding request static void newRequestSpeculativeDecoding(SizeType32 batchIdx, runtime::decoder_batch::Request const& request, diff --git a/cpp/include/tensorrt_llm/runtime/request.h b/cpp/include/tensorrt_llm/runtime/request.h index 983d19ad8c98..e8f851b7d777 100644 --- a/cpp/include/tensorrt_llm/runtime/request.h +++ b/cpp/include/tensorrt_llm/runtime/request.h @@ -31,9 +31,8 @@ class Request using TensorPtr = ITensor::SharedPtr; using BufferPtr = IBuffer::SharedPtr; - explicit Request(SizeType32 inputLen, std::optional maxNewTokens = std::nullopt) + explicit Request(SizeType32 inputLen) : inputLen(inputLen) - , maxNewTokens{maxNewTokens} { } @@ -41,7 +40,6 @@ class Request SizeType32 inputLen; // Input length without draft tokens, increasing with generation steps // optional parameters - std::optional maxNewTokens; // maximum number of tokens to generate for this request SizeType32 generatedTokensPerEngineStep{1}; // //! Optional parameters for speculative decoding diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 5bd20aade108..6350e6414618 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -302,26 +302,12 @@ void initializeLogProbs(DecodingOutput& dJointOutput, SizeType32 batchSlot, Samp void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, - runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream, SizeType32 maxSequenceLength) + runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(batchSlot >= 0); - BufferManager manager{std::make_shared(decoderStream.get())}; - auto const batchSize = decoderState.getMaxNumSequences(); - TLLM_CHECK(0 <= batchSize && batchSlot < batchSize); - - auto const inputLength = request.inputLen; - auto const numDecodingEngineTokens = request.generatedTokensPerEngineStep; - - // input - auto& dJointInput = decoderState.getJointDecodingInput(); - - initializeInputLengths(dJointInput, batchSlot, inputLength, request.maxNewTokens, - request.generatedTokensPerEngineStep, maxSequenceLength, manager); - // output auto& dJointOutput = decoderState.getJointDecodingOutput(); @@ -605,6 +591,8 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon TLLM_CHECK(llmReq->mSeqSlot.has_value()); auto const batchSlot = llmReq->mSeqSlot.value(); + auto const batchSize = decoderState.getMaxNumSequences(); + TLLM_CHECK(0 <= batchSlot && batchSlot < batchSize); auto const& samplingConfig = llmReq->mSamplingConfig; @@ -617,7 +605,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon auto const promptLen = llmReq->getPromptLen(); - auto decoderRequest = decoder_batch::Request{promptLen, llmReq->mMaxNewTokens}; + auto decoderRequest = decoder_batch::Request{promptLen}; if (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) { @@ -644,6 +632,11 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon decoderRequest.generatedTokensPerEngineStep = modelConfig.getMaxDecodingTokens(); } + auto const numDecodingEngineTokens = decoderRequest.generatedTokensPerEngineStep; + initializeInputLengths(dJointInput, batchSlot, promptLen, llmReq->mMaxNewTokens, numDecodingEngineTokens, + maxSequenceLength, decoderBufferManager); + decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); + initializeEmbeddingBias( dJointInput, batchSlot, llmReq->getEmbeddingBias(), logitsType, modelConfig, decoderBufferManager); @@ -653,8 +646,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon setupWords(dJointInput.stopWordsLists, llmReq->getStopWordsList(), dJointInput.stopWordsPtrs, dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, decoderBufferManager); - newRequest( - batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState, decoderStream, maxSequenceLength); + newRequest(batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState, decoderStream); auto& dJointOutput = decoderState.getJointDecodingOutput(); @@ -676,9 +668,6 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, decoderBufferManager); } - auto const numDecodingEngineTokens = decoderRequest.generatedTokensPerEngineStep; - decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); - // Speculative execution if (!decoderState.getSpeculativeDecodingMode().isNone()) { From f34b7e5606851e119402b01d8964fd605882c424 Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Fri, 1 Aug 2025 14:41:35 +0000 Subject: [PATCH 16/17] refactor: Rename newRequest to initializeOutputs - All other functionality has been moved out of the newRequest function. - Renamed newRequest to initializeOutputs to better reflect its purpose. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- .../batch_manager/createNewDecoderRequests.h | 5 ----- .../createNewDecoderRequests.cpp | 22 +++++++------------ 2 files changed, 8 insertions(+), 19 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index e129cf6e9ab6..394f7fb7bfa1 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -88,11 +88,6 @@ class CreateNewDecoderRequests : Algorithm SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const; private: - //! @brief Initialize the decoder at `batchSlot` with a new `request`. - static void newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, - SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, - runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream); - //! @brief Setups decoder internal tensors for new speculative decoding request static void newRequestSpeculativeDecoding(SizeType32 batchIdx, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 6350e6414618..5accd1c6a8d7 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -298,23 +298,15 @@ void initializeLogProbs(DecodingOutput& dJointOutput, SizeType32 batchSlot, Samp } } -} // namespace - -void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder_batch::Request const& request, - SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, - runtime::decoder::DecoderState& decoderState, CudaStream const& decoderStream) +void initializeOutputs(DecodingOutput& dJointOutput, SizeType32 batchSlot, SizeType32 maxDecodingEngineTokens, + BufferManager const& manager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - BufferManager manager{std::make_shared(decoderStream.get())}; - - // output - auto& dJointOutput = decoderState.getJointDecodingOutput(); - auto finishedSum = ITensor::slice(dJointOutput.finishedSum, batchSlot, 1); manager.setZero(*finishedSum); - for (SizeType32 ti = 0; ti < decoderState.getMaxDecodingEngineTokens(); ++ti) + for (SizeType32 ti = 0; ti < maxDecodingEngineTokens; ++ti) { TensorPtr const newTokensStepView = ITensor::slice(dJointOutput.newTokensSteps, ti, 1); newTokensStepView->squeeze(0); @@ -322,12 +314,14 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder manager.setZero(*newTokensVec); } - TensorPtr const finishedStepsSlice = ITensor::slice(decoderState.getFinishReasons(), batchSlot, 1); + TensorPtr const finishedStepsSlice = ITensor::slice(dJointOutput.finishReasons, batchSlot, 1); manager.setZero(*finishedStepsSlice); TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } +} // namespace + void CreateNewDecoderRequests::newRequestSpeculativeDecoding(SizeType32 batchIdx, runtime::decoder_batch::Request const& request, SamplingConfig const& samplingConfig, runtime::ModelConfig const& modelConfig, DecodingInput& jointDecodingInput, DecodingOutput& jointDecodingOutput, @@ -646,10 +640,10 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon setupWords(dJointInput.stopWordsLists, llmReq->getStopWordsList(), dJointInput.stopWordsPtrs, dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, decoderBufferManager); - newRequest(batchSlot, decoderRequest, samplingConfig, modelConfig, decoderState, decoderStream); - auto& dJointOutput = decoderState.getJointDecodingOutput(); + initializeOutputs(dJointOutput, batchSlot, decoderState.getMaxDecodingEngineTokens(), decoderBufferManager); + initializeLogProbs(dJointOutput, batchSlot, samplingConfig, decoderBufferManager); auto const& reqTokens = llmReq->getTokens(0); From b4d3a695b4df25dcada2f117c4f6287ad3088edc Mon Sep 17 00:00:00 2001 From: Robin Kobus <19427718+Funatiq@users.noreply.github.com> Date: Sat, 2 Aug 2025 14:11:00 +0000 Subject: [PATCH 17/17] refactor: Add setBeamWidth method to DecoderState - Introduced setBeamWidth method in DecoderState to allow setting the beam width for specific requests in a batch. - Updated CreateNewDecoderRequests to utilize the new setBeamWidth method, improving code clarity and maintainability. Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/decoderState.h | 5 +++++ cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp | 6 +++--- cpp/tensorrt_llm/runtime/decoderState.cpp | 5 +++++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/decoderState.h b/cpp/include/tensorrt_llm/runtime/decoderState.h index 8166156a0cc7..95d7ff0ffac9 100644 --- a/cpp/include/tensorrt_llm/runtime/decoderState.h +++ b/cpp/include/tensorrt_llm/runtime/decoderState.h @@ -173,6 +173,11 @@ class DecoderState //! @brief Workspace for beam search in streaming mode. [[nodiscard]] BeamSearchBuffers const& getBeamSearchBuffers() const; + //! @brief Set the beam width for a specific request in the batch. + //! @param batchIdx The index of the request in the batch. + //! @param beamWidth The beam width for the specified request. + void setBeamWidth(SizeType32 batchIdx, SizeType32 beamWidth); + //! @brief Cache indirection input for beam search. [[nodiscard]] TensorPtr getCacheIndirectionInput() const; diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 5accd1c6a8d7..16771709bb49 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -581,8 +581,6 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon { llmReq->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; - auto& dJointInput = decoderState.getJointDecodingInput(); - TLLM_CHECK(llmReq->mSeqSlot.has_value()); auto const batchSlot = llmReq->mSeqSlot.value(); auto const batchSize = decoderState.getMaxNumSequences(); @@ -595,7 +593,7 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon TLLM_CHECK_WITH_INFO(beamWidth <= maxBeamWidth, tc::fmtstr("Beam width (%d) must be smaller than maxBeamWidth (%d) passed to decoder setup function.", beamWidth, maxBeamWidth)); - dJointInput.beamWidths.at(batchSlot) = beamWidth; + decoderState.setBeamWidth(batchSlot, beamWidth); auto const promptLen = llmReq->getPromptLen(); @@ -626,6 +624,8 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon decoderRequest.generatedTokensPerEngineStep = modelConfig.getMaxDecodingTokens(); } + auto& dJointInput = decoderState.getJointDecodingInput(); + auto const numDecodingEngineTokens = decoderRequest.generatedTokensPerEngineStep; initializeInputLengths(dJointInput, batchSlot, promptLen, llmReq->mMaxNewTokens, numDecodingEngineTokens, maxSequenceLength, decoderBufferManager); diff --git a/cpp/tensorrt_llm/runtime/decoderState.cpp b/cpp/tensorrt_llm/runtime/decoderState.cpp index 19afac7fce75..abccbe60a137 100644 --- a/cpp/tensorrt_llm/runtime/decoderState.cpp +++ b/cpp/tensorrt_llm/runtime/decoderState.cpp @@ -644,6 +644,11 @@ void DecoderState::setGenerationSteps(std::vector const& generationS mJointDecodingInput->generationSteps = generationSteps; } +void DecoderState::setBeamWidth(SizeType32 batchIdx, SizeType32 beamWidth) +{ + mJointDecodingInput->beamWidths.at(batchIdx) = beamWidth; +} + DecodingInput& DecoderState::getJointDecodingInput() const { return *mJointDecodingInput;