diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index ce42493879e9..394f7fb7bfa1 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -75,27 +75,19 @@ 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; private: - //! @brief Initialize the decoder at `batchSlot` with a new `request`. Exposed only for static batching via - //! GptDecoderBatched::newBatch() - 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); - //! @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/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/include/tensorrt_llm/runtime/request.h b/cpp/include/tensorrt_llm/runtime/request.h index 1861ea843174..e8f851b7d777 100644 --- a/cpp/include/tensorrt_llm/runtime/request.h +++ b/cpp/include/tensorrt_llm/runtime/request.h @@ -31,26 +31,16 @@ 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) - , maxNewTokens{maxNewTokens} - , endId{endId} + explicit Request(SizeType32 inputLen) + : inputLen(inputLen) { } //! 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}; // - TensorPtr embeddingBias; // [vocabSizePadded], on gpu - 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 84fd13d5c18a..16771709bb49 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -44,14 +44,16 @@ namespace tensorrt_llm::batch_manager using SizeType32 = CreateNewDecoderRequests::SizeType32; using TensorPtr = CreateNewDecoderRequests::TensorPtr; +using SharedConstPtr = CreateNewDecoderRequests::SharedConstPtr; 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); @@ -79,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); } } @@ -127,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); @@ -141,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(); @@ -165,115 +167,122 @@ CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, ru std::move(lookaheadAlgoConfigs)}; } -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) +namespace { - 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 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& requestIds = request.ids; - auto const inputLength = request.inputLen; - auto const numDecodingEngineTokens = request.generatedTokensPerEngineStep; +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 - = request.maxNewTokens.value_or(maxSequenceLength - inputLength - numDecodingDraftEngineTokens); + 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)); - TLLM_CHECK(requestIds->getDataType() == TRTDataType::value); - auto const endId = request.endId.value_or(-1); - // input - auto& dJointInput = decoderState.getJointDecodingInput(); + TensorPtr const sequenceLimitLength{ + ITensor::slice(constPointerCast(dJointInput.sequenceLimitLength), batchSlot, 1)}; + runtime::kernels::invokeFill(*sequenceLimitLength, inputLength + maxNewTokens, manager.getStream()); - dJointInput.beamWidths.at(batchSlot) = beamWidth; - decoderState.setNumDecodingEngineTokens(batchSlot, numDecodingEngineTokens); + 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) +{ TensorPtr const endIdTensorPtr{ITensor::slice(constPointerCast(dJointInput.endIds), batchSlot, 1)}; - runtime::kernels::invokeFill(*endIdTensorPtr, endId, decoderStream); + 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) +{ TensorPtr const embeddingBiasSlice = ITensor::slice(constPointerCast(dJointInput.embeddingBias), batchSlot, 1); - if (request.embeddingBias) + if (embeddingBias.has_value()) { - 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(), + 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(), "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); + manager.copy(*embeddingBiasTensor, *embeddingBiasSlice); } else { manager.setZero(*embeddingBiasSlice); } +} - auto 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.has_value()) { - 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); - - 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); - - TensorPtr const inputLengths{ITensor::slice(constPointerCast(dJointInput.lengths), batchSlot, 1)}; - runtime::kernels::invokeFill(*inputLengths, inputLength, decoderStream); - - // output - auto& dJointOutput = decoderState.getJointDecodingOutput(); - auto const outputIdsShape = ITensor::makeShape({1, beamWidth, maxSequenceLength}); - - auto finishedSum = ITensor::slice(dJointOutput.finishedSum, batchSlot, 1); - manager.setZero(*finishedSum); - - for (SizeType32 ti = 0; ti < decoderState.getMaxDecodingEngineTokens(); ++ti) + // 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(*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] = wordsList; + } + else { - TensorPtr const newTokensStepView = ITensor::slice(dJointOutput.newTokensSteps, ti, 1); - newTokensStepView->squeeze(0); - auto newTokensVec = ITensor::slice(newTokensStepView, batchSlot, 1); - manager.setZero(*newTokensVec); + runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = 0; } +}; - TensorPtr const finishedStepsSlice = ITensor::slice(decoderState.getFinishReasons(), batchSlot, 1); - manager.setZero(*finishedStepsSlice); +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) @@ -287,49 +296,32 @@ void CreateNewDecoderRequests::newRequest(SizeType32 batchSlot, runtime::decoder auto logProbs = ITensor::slice(dJointOutput.logProbs, batchSlot, 1); manager.setZero(*logProbs); } +} - 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); +void initializeOutputs(DecodingOutput& dJointOutput, SizeType32 batchSlot, SizeType32 maxDecodingEngineTokens, + BufferManager const& manager) +{ + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto beamHypotheses = dJointOutput.beamHypotheses.slice(batchSlot, 1); - beamHypotheses.init(manager, endId); - } + auto finishedSum = ITensor::slice(dJointOutput.finishedSum, batchSlot, 1); + manager.setZero(*finishedSum); - // Speculative execution - if (numDecodingEngineTokens > 1 || decoderState.getSpeculativeDecodingMode().isDraftTokensExternal()) + for (SizeType32 ti = 0; ti < maxDecodingEngineTokens; ++ti) { - TLLM_CHECK(beamWidth == 1); - newRequestSpeculativeDecoding(batchSlot, request, samplingConfig, modelConfig, - decoderState.getJointDecodingInput(), decoderState.getJointDecodingOutput(), runtimeStream, decoderStream, - decoderState.getSpeculativeDecodingMode(), decoderState.getMaxDecodingEngineTokens()); + TensorPtr const newTokensStepView = ITensor::slice(dJointOutput.newTokensSteps, ti, 1); + newTokensStepView->squeeze(0); + auto newTokensVec = ITensor::slice(newTokensStepView, batchSlot, 1); + manager.setZero(*newTokensVec); } - // 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); + 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, @@ -557,11 +549,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 decoderBufferManager = BufferManager{std::make_shared(decoderStream.get())}; + unsigned decoderInputSize{0}; for (auto const& llmReq : finishedContextRequests) { @@ -586,26 +579,38 @@ CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedCon SizeType32 inputOffset{0}; for (auto const& llmReq : finishedContextRequests) { + llmReq->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; + + 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; + + 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)); + decoderState.setBeamWidth(batchSlot, 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->mSamplingConfig.normalizeLogProbs = mIsNormalizeLogProbs; if (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) { 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; } @@ -618,48 +623,77 @@ 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()) - { - decoderRequest.embeddingBias = getEmbeddingBias(logitsType, llmReq->getEmbeddingBias().value()); - } - if (llmReq->getBadWordsList().has_value()) + auto& dJointInput = decoderState.getJointDecodingInput(); + + 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); + + setupWords(dJointInput.badWordsLists, llmReq->getBadWordsList(), dJointInput.badWordsPtrs, + dJointInput.badWordsLens, dJointInput.maxBadWordsLen, batchSlot, decoderBufferManager); + + setupWords(dJointInput.stopWordsLists, llmReq->getStopWordsList(), dJointInput.stopWordsPtrs, + dJointInput.stopWordsLens, dJointInput.maxStopWordsLen, batchSlot, decoderBufferManager); + + auto& dJointOutput = decoderState.getJointDecodingOutput(); + + initializeOutputs(dJointOutput, batchSlot, decoderState.getMaxDecodingEngineTokens(), decoderBufferManager); + + 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) + decoderBufferManager.copy(reqTokens.data(), *requestIds); + auto const endId = llmReq->mEndId.value_or(-1); + + initializeRequestIds(dJointInput, dJointOutput, batchSlot, requestIds, endId, beamWidth, maxSequenceLength, + decoderBufferManager); + + if (beamWidth > 1) { - // 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); + initializeBeamSearch( + dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, decoderBufferManager); } - if (llmReq->getStopWordsList().has_value()) + + // Speculative execution + if (!decoderState.getSpeculativeDecodingMode().isNone()) { - decoderRequest.stopWordsList - = bufferManager.copyFrom(*llmReq->getStopWordsList().value(), MemoryType::kGPU); - decoderRequest.stopWordsList->squeeze(0); - } + TLLM_CHECK(beamWidth == 1); - TLLM_CHECK(llmReq->mSeqSlot.has_value()); - newRequest(llmReq->mSeqSlot.value(), decoderRequest, llmReq->mSamplingConfig, modelConfig, decoderState, - runtimeStream, decoderStream, maxSequenceLength); + 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(requestIds); + + 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()); + } decoderRequests.push_back(decoderRequest); 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/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; 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,