diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4dcede1c4751..fcd6cdcfc09a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -306,7 +306,6 @@ common-files: &common_files | tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py | tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py | tensorrt_llm/_torch/pyexecutor/llm_request.py | - tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py | tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | tensorrt_llm/_torch/pyexecutor/model_engine.py | tensorrt_llm/_torch/pyexecutor/model_loader.py | @@ -619,7 +618,6 @@ common-files: &common_files | tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_best_of_n.py | - tests/unittest/_torch/sampler/test_trtllm_sampler.py | tests/unittest/_torch/speculative/test_eagle3.py | tests/unittest/_torch/test_connector.py | tests/unittest/_torch/test_torch_multi_arange.py | @@ -1072,7 +1070,6 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py | tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py | tensorrt_llm/_torch/pyexecutor/llm_request.py | - tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py | tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | tensorrt_llm/_torch/pyexecutor/model_engine.py | tensorrt_llm/_torch/pyexecutor/model_loader.py | @@ -1385,7 +1382,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_best_of_n.py | - tests/unittest/_torch/sampler/test_trtllm_sampler.py | tests/unittest/_torch/speculative/test_eagle3.py | tests/unittest/_torch/test_connector.py | tests/unittest/_torch/test_torch_multi_arange.py | diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h deleted file mode 100644 index 600927af9645..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ /dev/null @@ -1,88 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::runtime -{ -class SamplingConfig; - -namespace decoder -{ -class DecoderState; -} // namespace decoder - -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ -class MedusaBuffers; -class DecoderInputBuffers; - -class CreateNewDecoderRequests : Algorithm -{ -public: - constexpr static auto name{"CreateNewDecoderRequests"}; - - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using SamplingConfig = tensorrt_llm::runtime::SamplingConfig; - using CudaStream = tensorrt_llm::runtime::CudaStream; - using TensorPtr = runtime::ITensor::SharedPtr; - using SharedConstPtr = runtime::ITensor::SharedConstPtr; - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - CreateNewDecoderRequests(bool speculativeDecodingFastLogits, bool isLeaderInOrchMode, bool isNormalizeLogProbs) - : mSpeculativeDecodingFastLogits(speculativeDecodingFastLogits) - , mIsLeaderInOrchMode(isLeaderInOrchMode) - , mIsNormalizeLogProbs(isNormalizeLogProbs) - { - } - - [[nodiscard]] std::tuple, std::vector, - std::vector> - operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - tensorrt_llm::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, - tensorrt_llm::DataType logitsType, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, runtime::CudaStream const& runtimeStream, - runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, - OptionalRef medusaBuffers) const; - -private: - bool mSpeculativeDecodingFastLogits; - bool mIsLeaderInOrchMode; - bool mIsNormalizeLogProbs; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h b/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h deleted file mode 100644 index df507cf10013..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h +++ /dev/null @@ -1,186 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -namespace tensorrt_llm::runtime::decoder -{ -class DecoderState; -} - -namespace tensorrt_llm::batch_manager -{ - -class DecoderInputBuffers -{ -public: - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorConstPtr = runtime::ITensor::SharedConstPtr; - - explicit DecoderInputBuffers( - SizeType32 maxBatchSize, SizeType32 maxDecoderSteps, runtime::BufferManager const& manager); - - void setupMedusaLogits(SizeType32 maxNumSequences, runtime::ModelConfig const& modelConfig); - - //! Buffers for decoder setup - - //! Input IDs of new requests, [maxBatchSize] - TensorPtr inputsIds; - //! Batch slots for setup step, [maxBatchSize] - TensorPtr setupBatchSlots; - TensorPtr setupBatchSlotsDevice; - //! Helper buffer for copying sequence lengths, [maxBatchSize] - TensorPtr fillValues; - TensorPtr fillValuesDevice; - - //! Buffers for decoder forward - - //! Requests for considered in decoder forward - RequestVector decoderRequests; - - //! Logits of decoder requests - std::vector decoderLogits; - - //! Maximum number of decoding steps of decoder requests. - //! This is only more than 1 for external draft tokens speculative decoding. - SizeType32 maxDecoderSteps{1}; - - //! Batch slots for all decoder steps, [maxDecoderSteps][maxBatchSize] - std::vector forwardBatchSlots; - - //! Logits for requests in forwardBatchSlots (in the same order). - //! [maxDecoderSteps][batchSize][1, beamWidth, vocabSizePadded], on gpu - std::vector> batchLogits; - - //! Logits for speculative decoding (Medusa). - //! The vector is sparse, only slots in forwardBatchSlots are used. - //! [maxBatchSize][maxAcceptedDraftTokensPerStep][maxDraftTokens + 1, vocabSizePadded] - std::vector> predictedDraftLogits; -}; - -class DecoderOutputBuffers -{ -public: - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - - DecoderOutputBuffers(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxSeqLen, - SizeType32 maxTokensPerStep, runtime::BufferManager const& manager); - - void enableLookaheadDecoding(SizeType32 maxNumSequences, SizeType32 maxTokensPerStep); - void disableLookaheadDecoding(SizeType32 maxNumSequences); - - void setupSpeculativeDecoding( - SizeType32 maxNumSequences, SizeType32 maxTokensPerStep, runtime::ModelConfig const& modelConfig); - - TensorPtr sequenceLengthsHost; // [mMaxNumRequests, beamWidth], pinned host tensor - TensorPtr newOutputTokensHost; // [maxTokensPerStep, mMaxNumRequests, beamWidth] - TensorPtr cumLogProbsHost; // [mMaxNumRequests, beamWidth] - TensorPtr logProbsHost; // [mMaxNumRequests, beamWidth, maxSeqLen] - TensorPtr finishedSumHost; // [mMaxNumRequests], pinned host tensor - TensorPtr finishReasonsHost; // [mMaxNumRequests, beamWidth], pinned host tensor - - // speculative decoding buffers - TensorPtr nextDraftTokensHost; // [mMaxNumRequests, maxTokensPerStep-1] - TensorPtr prevDraftTokensLengthsHost; // [mMaxNumRequests] - TensorPtr nextDraftTokensLengthsHost; // [mMaxNumRequests] -}; - -class DecoderStepAsyncSend -{ -public: - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - - DecoderStepAsyncSend(DecoderOutputBuffers const& decoderOutputBuffers, - runtime::decoder::DecoderState const& decoderState, bool returnLogProbs, SizeType32 maxBeamWidth, - bool useMedusa, mpi::MpiComm const& commSession, int peer); - - ~DecoderStepAsyncSend(); - - static void recv(DecoderOutputBuffers const& decoderOutputBuffers, - runtime::decoder::DecoderState const& decoderState, bool returnLogProbs, SizeType32 maxBeamWidth, - bool useMedusa, mpi::MpiComm const& commSession, int peer); - - static void bcast(DecoderOutputBuffers const& decoderOutputBuffers, - runtime::decoder::DecoderState const& decoderState, bool returnLogProbs, SizeType32 maxBeamWidth, - bool useMedusa, mpi::MpiComm const& commSession, int root); - -private: - std::unique_ptr mRequest1; - std::unique_ptr mRequest2; - std::unique_ptr mRequest3; - std::unique_ptr mRequest4; - std::unique_ptr mRequest5; - std::unique_ptr mRequest6; - std::unique_ptr mRequest7; - std::unique_ptr mRequest8; - std::unique_ptr mRequest9; -}; - -class SlotDecoderBuffers -{ -public: - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - - TensorPtr outputIds; // [beamWidth, maxSeqLen], outputIds of single batch slot - TensorPtr outputIdsHost; // [beamWidth, maxSeqLen], outputIds of single batch slot - TensorPtr sequenceLengths; // [beamWidth] - TensorPtr sequenceLengthsHost; // [beamWidth] - TensorPtr cumLogProbs; // [beamWidth] - TensorPtr cumLogProbsHost; // [beamWidth] - TensorPtr logProbs; // [beamWidth, maxSeqLen] - TensorPtr logProbsHost; // [beamWidth, maxSeqLen] - TensorPtr finishReasonsHost; // [beamWidth] - - SlotDecoderBuffers(SizeType32 maxBeamWidth, SizeType32 maxSeqLen, runtime::BufferManager const& manager); -}; - -class DecoderSlotAsyncSend -{ -public: - using TensorPtr = runtime::ITensor::SharedPtr; - - DecoderSlotAsyncSend(TensorPtr const& outputIds, TensorPtr const& sequenceLengths, TensorPtr const& cumLogProbs, - TensorPtr const& logProbs, bool returnLogProbs, mpi::MpiComm const& commSession, int peer); - - DecoderSlotAsyncSend( - SlotDecoderBuffers const& slotDecoderBuffers, bool returnLogProbs, mpi::MpiComm const& commSession, int peer); - - ~DecoderSlotAsyncSend(); - - static void recv( - SlotDecoderBuffers const& slotDecoderBuffers, bool returnLogProbs, mpi::MpiComm const& commSession, int peer); - -private: - std::unique_ptr mRequest1; - std::unique_ptr mRequest2; - std::unique_ptr mRequest3; - std::unique_ptr mRequest4; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h b/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h deleted file mode 100644 index 5342591840a8..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h +++ /dev/null @@ -1,66 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/promptTuningParams.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::batch_manager -{ - -class MedusaBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using ITensor = tensorrt_llm::runtime::ITensor; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep); - - void insertInputTensors( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::WorldConfig const& worldConfig) const; - -public: - TensorPtr medusaLogitsDevice; // [maxAcceptedDraftTokens, maxBatchSize, maxDraftTokens + 1, vocabSizePadded], on gpu - - TensorPtr attentionPackedMaskDevice; // [maxBatchSize, maxDraftTokens + 1, numPackedMasks], on gpu - TensorPtr attentionPackedMaskHost; // [maxBatchSize, maxDraftTokens + 1, numPackedMasks], on pinned - - TensorPtr medusaGenerationLengthsDevice; // [maxBatchSize], on gpu - TensorPtr medusaGenerationLengthsHost; // [maxBatchSize], on pinned - - TensorPtr medusaPositionOffsetsDevice; // [maxBatchSize, maxDraftTokens + 1], on gpu - TensorPtr medusaPositionOffsetsHost; // [maxBatchSize, maxDraftTokens + 1], on pinned - - TensorPtr medusaTreeIdsDevice; // [maxBatchSize, maxDraftTokens + 1], on gpu - TensorPtr medusaTreeIdsHost; // [maxBatchSize, maxDraftTokens + 1], on pinned - - TensorPtr medusaPathsDevice; // [maxBatchSize, maxDraftTokens + 1, maxAcceptedDraftTokens + 1], on gpu - TensorPtr medusaPathsHost; // [maxBatchSize, maxDraftTokens + 1, maxAcceptedDraftTokens + 1], on pinned - - TensorPtr medusaUseSpecDecoding; // [1], on cpu - - std::vector mTopKs; // [maxAcceptedDraftTokens] -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index acc0efe18966..4fc8d30e2402 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -63,6 +63,11 @@ class DataTransceiverState; class SamplingConfig { public: + /// @brief Largest beam width a request may ask for. + static constexpr SizeType32 kMaxBeamWidth = 1024; + /// @brief Largest variable-beam-width schedule (beamWidthArray) a request may supply. + static constexpr SizeType32 kMaxBeamWidthArrayLength = 8; + /// @brief Constructor for SamplingConfig /// See description of parameters below explicit SamplingConfig(SizeType32 beamWidth = 1, std::optional const& topK = std::nullopt, diff --git a/cpp/include/tensorrt_llm/runtime/decoderState.h b/cpp/include/tensorrt_llm/runtime/decoderState.h deleted file mode 100644 index ea2c767c0478..000000000000 --- a/cpp/include/tensorrt_llm/runtime/decoderState.h +++ /dev/null @@ -1,244 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "decodingInput.h" -#include "decodingOutput.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/speculativeDecodingMode.h" - -namespace tensorrt_llm::runtime::decoder -{ - -class BeamSearchBuffers -{ -public: - explicit BeamSearchBuffers(BufferManager const& bufferManager); - - void reshape(SizeType32 maxBeamWidth, SizeType32 maxSequenceLength); - - // temporary buffers for the beam search + streaming case - DecodingOutput::BeamHypotheses mOutputBeamHypotheses; - // will store a slice of DecodingOutput::cumLogProbs - DecodingOutput::TensorPtr mCumLogProbsTmp; - SizeType32 mNumSMs; -}; - -class DecoderState -{ -public: - using TensorPtr = ITensor::SharedPtr; - using LlmRequestPtr = std::shared_ptr; - using RequestVector = std::vector; - using DecodingInputPtr = std::unique_ptr; - using DecodingOutputPtr = std::unique_ptr; - - DecoderState(); - - //! @brief Setup buffers for the decoder excluding speculative decoding. - void setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, tensorrt_llm::DataType dtype, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); - - //! @brief Setup buffers for the cache indirection. - //! @details This is used for beam search on pipeline parallel ranks without a decoder. - void setupCacheIndirection(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - BufferManager const& bufferManager); - - //! @brief Setup buffers for speculative decoding. - void setupSpeculativeDecoding(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, BufferManager const& bufferManager); - - //! @brief Disable lookahead decoding. - void disableLookahead(RequestVector const& genRequests); - - //! @returns [batchSize], number of finished sequences per request, on gpu - [[nodiscard]] TensorPtr getFinishedSum() const; - - //! @returns [batchSize, beamWidth], finished states of type FinishedState, on gpu - [[nodiscard]] TensorPtr getFinishReasons() const; - - //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token - //! ids without padding, on gpu. In case of beam search, contains the ungathered data. - [[nodiscard]] TensorPtr getIds() const; - - //! @param batchIdx index of the batch - //! @returns [maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without - //! padding for request `batchIdx`, on gpu. In case of beam search, contains the ungathered data. - [[nodiscard]] TensorPtr getIds(SizeType32 batchIdx) const; - - //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], only used for beam search. It contains - //! gathered token ids without padding, on gpu. - [[nodiscard]] TensorPtr getGatheredIds() const; - - //! @param batchIdx index of the batch - //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], only used for beam search. It contains - //! gathered token ids without padding for request `batchIdx`, on gpu. - [[nodiscard]] TensorPtr getGatheredIds(SizeType32 batchIdx) const; - - //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains parent ids collected during beam - //! search without padding, on gpu - [[nodiscard]] TensorPtr getParentIds() const; - - //! @returns [batchSize, maxBeamWidth], cumulative log probabilities (per beam), on gpu - [[nodiscard]] TensorPtr getCumLogProbs() const; - - //! @returns [maxBeamWidth], cumulative log probabilities (per beam), on gpu - [[nodiscard]] TensorPtr getCumLogProbs(SizeType32 batchIdx) const; - - //! @returns [batchSize, maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu - [[nodiscard]] TensorPtr getLogProbs() const; - - //! @returns [maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu - [[nodiscard]] TensorPtr getLogProbs(SizeType32 batchIdx) const; - - //! @returns [batchSize, maxBeamWidth], sequence lengths, on gpu - [[nodiscard]] TensorPtr getSequenceLengths() const; - - //! @param batchIdx index of the batch - //! @returns [maxBeamWidth], sequence lengths for request `batchIdx`, on gpu - [[nodiscard]] TensorPtr getSequenceLengths(SizeType32 batchIdx) const; - - //! @brief Get maxTokensPerStep tokens generated in the last forward pass - //! @returns [maxTokensPerStep, batchSize, maxBeamWidth], tokens generated in last forward pass, on gpu - [[nodiscard]] TensorPtr getAllNewTokens() const; - - //! @returns [batchSize, maxDraftTokens], predicted draft tokens for next step, on gpu - [[nodiscard]] TensorPtr getNextDraftTokens() const; - - //! @returns [batchSize], predicted draft tokens lengths for previous step, on gpu - [[nodiscard]] TensorPtr getPrevDraftTokensLengths() const; - - //! @returns [batchSize], predicted draft tokens lengths for next step, on gpu - [[nodiscard]] TensorPtr getNextDraftTokensLengths() const; - - //! @returns [batchSize + 1], exclusive sum of accepted draft token lengths, on gpu - [[nodiscard]] TensorPtr getAcceptedLengthsCumSum() const; - - //! @returns [batchSize, maxAcceptedDraftTokensPerStep], accepted paths packed into continuous tensor, on gpu - [[nodiscard]] TensorPtr getAcceptedPackedPaths() const; - - [[nodiscard]] SizeType32 getMaxNumSequences() const; - - [[nodiscard]] SizeType32 getMaxBeamWidth() const; - - [[nodiscard]] SizeType32 getMaxSequenceLength() const; - - [[nodiscard]] SizeType32 getMaxDecodingDecoderTokens() const; - - [[nodiscard]] SizeType32 getMaxDecodingEngineTokens() const; - - //! @brief Get the number of tokens for all requests in the batch. - //! @returns The number of tokens for all requests in the batch. - [[nodiscard]] std::vector const& getNumDecodingEngineTokens() const; - - //! @brief Get the number of tokens for a specific request in the batch. - //! @param batchIdx The index of the request in the batch. - //! @returns The number of tokens for the specified request. - [[nodiscard]] SizeType32 getNumDecodingEngineTokens(SizeType32 batchIdx) const; - - //! @brief Set the number of tokens for a specific request in the batch. - //! @param batchIdx The index of the request in the batch. - //! @param numTokens The number of tokens for the specified request. - void setNumDecodingEngineTokens(SizeType32 batchIdx, SizeType32 numTokens); - - //! @brief Get the speculative decoding mode. - [[nodiscard]] SpeculativeDecodingMode getSpeculativeDecodingMode() const; - - //! @brief Get the explicit draft tokens buffers. - [[nodiscard]] ExplicitDraftTokensBuffers::Inputs const& getExplicitDraftTokensBuffers() const; - - //! @brief Get the eagle buffers. - [[nodiscard]] EagleBuffers::Inputs const& getEagleBuffers() const; - - //! @brief Get the lookahead buffers. - [[nodiscard]] LookaheadDecodingBuffers const& getLookaheadBuffers() const; - - //! @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; - - //! @brief Cache indirection output for beam search. - [[nodiscard]] TensorPtr getCacheIndirectionOutput() const; - - //! @brief Get the generation steps for all requests in the batch. - //! @returns The generation steps for all requests in the batch. - [[nodiscard]] std::optional> const& getGenerationSteps() const; - - //! @brief Set the generation steps for all requests in the batch. - //! @param generationSteps The generation steps for all requests in the batch. - void setGenerationSteps(std::vector const& generationSteps); - - //! @brief Stateful inputs for the decoder. Allocated for maxNumSequences slots. - [[nodiscard]] DecodingInput& getJointDecodingInput() const; - - //! @brief Stateful outputs for the decoder. Allocated for maxNumSequences slots. - [[nodiscard]] DecodingOutput& getJointDecodingOutput() const; - -private: - void setupBuffers(tensorrt_llm::DataType dtype, BufferManager const& bufferManager); - void reshapeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, BufferManager const& bufferManager); - - void setupCacheIndirectionBuffers(BufferManager const& bufferManager); - void reshapeCacheIndirectionBuffers( - SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow); - - void setupSpeculativeDecodingBuffers(SpeculativeDecodingMode speculativeDecodingMode, tensorrt_llm::DataType dtype, - BufferManager const& bufferManager); - void reshapeSpeculativeDecodingBuffers(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - BufferManager const& bufferManager); - - SizeType32 mMaxNumSequences{}; - SizeType32 mMaxBeamWidth{}; - SizeType32 mMaxSequenceLength{}; - - //! @brief Stateful inputs for the decoder. Allocated for maxNumSequences slots. - DecodingInputPtr mJointDecodingInput; - //! @brief Stateful outputs for the decoder. Allocated for maxNumSequences slots. - DecodingOutputPtr mJointDecodingOutput; - - //! @brief Workspace for beam search in streaming mode. - std::unique_ptr mBeamSearchBuffers; - - // How many tokens for one request can be processed per mDecoders call. - // It is maxDecodingTokens for non speculative decoding and Draft model approach. - // Otherwise it is 1. - SizeType32 mMaxDecodingDecoderTokens{1}; - // How many tokens predicted by the engine for one request. - // It is maxDecodingTokens. >= 1 for speculative decoding and == 1 for non speculative decoding. - SizeType32 mMaxDecodingEngineTokens{1}; - - //! @brief [batchSize], the num tokens of each request. - std::vector mNumDecodingEngineTokens; - - SpeculativeDecodingMode mSpeculativeDecodingMode{SpeculativeDecodingMode::None()}; -}; - -} // namespace tensorrt_llm::runtime::decoder diff --git a/cpp/include/tensorrt_llm/runtime/decodingInput.h b/cpp/include/tensorrt_llm/runtime/decodingInput.h deleted file mode 100644 index 4344f423ac11..000000000000 --- a/cpp/include/tensorrt_llm/runtime/decodingInput.h +++ /dev/null @@ -1,171 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -#include -#include - -namespace tensorrt_llm::runtime -{ - -/// @brief Represents the inputs to the decoder. -/// @details This input type is assumed immutable. It represents whatever the decoder received initially, and can always -/// be referred to as such. -class DecodingInput -{ -public: - using TensorConstPtr = ITensor::SharedConstPtr; - using TensorPtr = ITensor::SharedPtr; - - DecodingInput() = default; - - //! Mandatory parameters - //! The index of the decoding step we are on. Only used in Python runtime - SizeType32 step{}; - //! The maximum number of tokens to decode - SizeType32 maxLength{}; - //! The maximum length of the attention window to consider while decoding - SizeType32 maxAttentionWindow{}; - //! The number of tokens to use as attention sinks, https://arxiv.org/html/2309.17453v3 - SizeType32 sinkTokenLength{}; - //! The number of samples in the batch - SizeType32 batchSize{}; - //! The beam widths of each request, [batchSize] - std::vector beamWidths; - //! The maximum value in the `stopWordsLens` tensor - SizeType32 maxStopWordsLen{}; - //! The maximum value in the `badWordsLens` tensor - SizeType32 maxBadWordsLen{}; - //! The output of the model forward computation, a probability distribution over the vocabulary - //! [batchSize][numGenTokens, beamWidth, vocabSizePadded] on gpu - std::vector logitsVec; - //! The end ids, [batchSize * beamWidth] on gpu - TensorConstPtr endIds; - //! Address map of the linear batch id to to the seq slots, [batchSize] on pinned, int32_t - TensorConstPtr batchSlots; - - //! Optional parameters - //! Finished states at current iteration (skip decoding step of a request if true), [batchSize, beamWidth] on gpu - TensorConstPtr finishReasons; - //! The maximum sequence length for each sequence in the batch, [batchSize] on gpu - TensorConstPtr sequenceLimitLength; - TensorConstPtr embeddingBias; // [batchSize, vocabSizePadded] on gpu - TensorConstPtr lengths; // [batchSize, beamWidth] on gpu - std::vector badWordsLists; // [batchSize][2, badWordsLength] on gpu - TensorConstPtr badWordsPtrs; // [batchSize][2, badWordsLength] on pinned - TensorConstPtr badWordsLens; // [batchSize] on gpu - std::vector stopWordsLists; // [batchSize][2, stopWordsLength] on gpu - TensorConstPtr stopWordsPtrs; // [batchSize][2, stopWordsLength] on pinned - TensorConstPtr stopWordsLens; // [batchSize] on pinned - TensorConstPtr noRepeatNgramSize; // [batchSize] on gpu - - //! Parameters for beam search - //! KV cache index for beam search, [batchSize, beamWidth, maxSeqLen] on gpu - TensorPtr cacheIndirection; - //! Steps of each request, for Variable-Beam-Width-Search, [batchSize] - std::optional> generationSteps; - - // Medusa - class MedusaInputs - { - public: - //! [batchSize, maxTokensPerStep, maxMedusaHeads + 1], on gpu - TensorConstPtr medusaPaths; - //! [batchSize, maxTokensPerStep], on gpu - TensorConstPtr medusaTreeIds; - //! [batchSize][maxAcceptedDraftTokensPerStep][maxDraftTokens + 1, vocabSizePadded], on gpu - std::vector> medusaLogits; - //! [batchSize], on gpu - TensorPtr medusaCurTokensPerStep; - //! [batchSize], on gpu - TensorConstPtr medusaTargetTokensPerStep; - }; - - class ExternalDraftTokensInputs - { - public: - TensorPtr draftLogits; - TensorPtr draftLogitsHost; - TensorPtr draftProbs; - TensorPtr targetProbs; - TensorPtr numDraftTokens; - TensorPtr numDraftTokensHost; - TensorPtr draftTokenIds; - TensorPtr draftTokenIdsHost; - TensorPtr useDraftLogits; - TensorPtr useDraftLogitsHost; - - SizeType32 step; - float constantThreshold; - bool useRandomAcceptanceThreshold; - }; - - class ExplicitDraftTokensInputs - { - public: - TensorConstPtr nextDraftTokens; // [batchSize, maxNumPaths, maxPathLen] - TensorConstPtr nextFlatTokens; // [batchSize * maxDecodingTokens] - TensorConstPtr nextDraftIndices; // [batchSize, maxNumPaths, maxPathLen] - TensorConstPtr nextDraftProbs; // [batchSize, maxNumPaths, maxDraftPathLen, vocabSize] - TensorConstPtr lastDraftTokens; // [batchSize, maxNumPaths, maxPathLen] - TensorConstPtr lastDraftIndices; // [batchSize, maxNumPaths, maxPathLen] - TensorConstPtr masks; // [batchSize, maxDecodingTokens, maxDecodingTokens], bool - TensorConstPtr packedPositionIds; // [batchSize * maxDecodingTokens] - TensorConstPtr bestPathLengths; // [batchSize] - TensorConstPtr bestPathIndices; // [batchSize] - TensorConstPtr nextGenerationLengths; // [batchSize] - TensorConstPtr lastPositionIdsBase; // [batchSize] - TensorConstPtr lastGenerationLengths; // [batchSize] - TensorConstPtr maxGenLengthDevice; // [1] - TensorConstPtr seqSlots; // [batchSize] - }; - - struct LookaheadInputs - { - TensorPtr tokensPerStep; - }; - - struct EagleInputs - { - TensorConstPtr nextDraftTokens; // [batchSize, maxDecodingDraftTokens] - TensorConstPtr nextDraftLens; // [batchSize] - TensorConstPtr nextDraftPaths; // [batchSize, maxDecodingTokens, maxPathLen] - TensorConstPtr lastDraftTokens; // [batchSize, maxNumPaths, maxPathLen] - TensorConstPtr lastDraftLens; // [batchSize] - TensorConstPtr lastDraftPaths; // [batchSize, maxDecodingTokens, maxPathLen] - TensorConstPtr acceptedTokens; // [batchSize, maxPathLen] - TensorConstPtr acceptedLens; // [batchSize] - TensorConstPtr acceptedPathIds; // [batchSize] - TensorConstPtr chunkedContextNextTokens; // [batchSize] - TensorConstPtr seqSlots; // [batchSize] - }; - - std::optional medusaInputs; - - std::optional explicitDraftTokensInputs; - - std::optional lookaheadInputs; - - std::optional externalDraftTokensInputs; - - std::optional eagleInputs; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/decodingOutput.h b/cpp/include/tensorrt_llm/runtime/decodingOutput.h deleted file mode 100644 index 55b56335d393..000000000000 --- a/cpp/include/tensorrt_llm/runtime/decodingOutput.h +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/eagleBuffers.h" -#include "tensorrt_llm/runtime/explicitDraftTokensBuffers.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/lookaheadBuffers.h" -#include -#include - -namespace tensorrt_llm::batch_manager -{ -class LookaheadDecodingBuffers; -} // namespace tensorrt_llm::batch_manager - -namespace tensorrt_llm::runtime -{ -class DecodingOutput -{ -public: - using TensorPtr = ITensor::SharedPtr; - - // BS: batch_size, BM: beam_width, MSL: max_seq_length - // All TensorPtr without special comments are on gpu - - class BeamHypotheses - { - public: - // Keep same as cpp/tensorrt_llm/kernels/beamSearchKernels.h - TensorPtr outputIdsCBA; // [BS, BM*2, MSL] - TensorPtr logProbsCBA; // [BS, BM*2, MSL] - TensorPtr sequenceLengthsCBA; // [BS, BM*2] - TensorPtr cumLogProbsCBA; // [BS, BM*2] - TensorPtr normedScoresCBA; // [BS, BM*2] - TensorPtr numBeamsCBA; // [BS] - TensorPtr minNormedScoresCBA; // [BS] - TensorPtr batchDones; // [BS] - - void empty(BufferManager const& manager); - - void reshape(SizeType32 batchSize, SizeType32 beamWidth, SizeType32 maxSequenceLength); - - void release(); - - void init(BufferManager const& manager, TokenIdType endId); - - BeamHypotheses slice(SizeType32 batchIndex, SizeType32 size) const; - }; - - static float constexpr kNegativeInfinity = -1e20f; - - DecodingOutput() = default; - - //! Mandatory parameters - //! Previously generated token ids for all steps before DecodingInput.step, [BS, BM, MSL] - TensorPtr ids; - //! The tokens computed during the gatherTree step, [BS, BM, MSL] - //! Necessary for "Streaming + Beam Search" mode since beam search kernels store ungathered tokens in `ids`. - TensorPtr gatheredIds; - //! New tokens at each generated token of maxTokensPerStep, [maxTokensPerStep, BS, BM] - TensorPtr newTokensSteps; - //! A view of newTokensSteps for the current token, [BS, BM] - TensorPtr newTokens; - //! A Vector of views on newTokensSteps for each token [BS, BM]. - std::vector newTokensVec; - - //! Optional parameters - //! FinishedState by decoding if any of the stop conditions are met or if DecodingInput.finished is true, [BS, BM] - TensorPtr finishReasons; - //! The sum of finished sequences per request, in pinned memory, [BS] - TensorPtr finishedSum; - - //! Mandatory parameters for Beam Search - //! log-probility of generated tokens, [BS, BM, MSL], float - TensorPtr logProbs; - //! Sum log-probility of all generated tokens, [BS, BM] - TensorPtr cumLogProbs; - //! Index of the beam where the previous token is, [BS, BM, MSL] - TensorPtr parentIds; - //! Total sequence lengths including padding, [BS, BM] - TensorPtr lengths; - //! K/V indirection for next generation step, [BS, BM, MSL] - TensorPtr cacheIndirection; - //! Buffer used to store the transpose of the logProbs, [MSL, BS, BM] - TensorPtr logProbsTiled; - - BeamHypotheses beamHypotheses; - - // Speculative decoding - class SpeculativeDecodingOutputs - { - public: - TensorPtr nextDraftTokens; // [maxBatchSize, maxDraftTokens] - TensorPtr nextDraftTokensLen; // [maxBatchSize] - TensorPtr prevDraftTokensLen; // [maxBatchSize] - TensorPtr acceptedTokensLen; // [maxBatchSize] - TensorPtr acceptedLengthsCumSum; // [maxBatchSize + 1] - TensorPtr pathsOffsets; // [maxBatchSize, maxAcceptedDraftTokensPerStep] - }; - - std::optional speculativeDecodingOutputs; - - std::optional explicitDraftTokensBuffers; - - std::optional lookaheadOutputs; - - std::optional eagleBuffers; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/eagleBuffers.h b/cpp/include/tensorrt_llm/runtime/eagleBuffers.h deleted file mode 100644 index d6754f68f4f5..000000000000 --- a/cpp/include/tensorrt_llm/runtime/eagleBuffers.h +++ /dev/null @@ -1,182 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/eagleModule.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -namespace tensorrt_llm::batch_manager -{ -class LlmRequest; -} - -namespace tensorrt_llm::runtime -{ - -class EagleBuffers -{ -public: - using LlmRequestPtr = std::shared_ptr; - using RequestVector = std::vector; - using SizeType32 = runtime::SizeType32; - using ITensor = runtime::ITensor; - using BufferPtr = runtime::IBuffer::SharedPtr; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - // The datastruct is used for runtime buffer that is holding runtime state per request (shape starts with - // maxBatchSize) and for engine inputs (shape starts with numSequences). - class Inputs - { - public: - //! [maxBatchSize] or [numSequences] - TensorPtr temperatures; - //! [maxBatchSize] or [numSequences] - TensorPtr posteriorAlpha; - //! [maxBatchSize] or [numSequences] - TensorPtr posteriorThreshold; - //! [maxBatchSize] or [numSequences] - TensorPtr randomDataSample; - //! [maxBatchSize, maxDecodingTokens] or [numSequences, maxDecodingTokens] - TensorPtr randomDataValidation; - //! [maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens] - TensorPtr draftTokens; - //! [maxBatchSize] or [numSequences] - TensorPtr draftLens; - //! [maxBatchSize, maxNumPaths, maxPathLen] - //! or [numSequences, maxNumPaths, maxPathLen] - TensorPtr draftPaths; - //! [maxBatchSize, maxNumPaths, maxPathLen] - //! or [numSequences, maxNumPaths, maxPathLen] - TensorPtr draftPathsHost; - //! [maxBatchSize] or [numGenSequences] - TensorPtr specDecodingGenerationLengths; - //! [maxBatchSize] or [numGenSequences] - TensorPtr specDecodingGenerationLengthsHost; - //! [maxBatchSize, maxDecodingTokens, ceil(maxDecodingTokens / 32)] - //! or [numGenSequences, maxDecodingTokens, ceil(maxDecodingTokens / 32)] - TensorPtr specDecodingPackedMasks; - //! [maxBatchSize] or [numGenSequences] - TensorPtr specDecodingPositionOffsets; - //! [maxBatchSize] or [numSequences] - TensorPtr eagleNetCtxRequestTypesHost; - //! [maxBatchSize] or [numSequences] - TensorPtr eagleNetCtxContextLengthsHost; - //! [maxBatchSize] or [numSequences] - TensorPtr eagleNetCtxPastKeyValueLengthsHost; - //! [maxBatchSize] or [numSequences] - TensorPtr eagleNetGenRequestTypesHost; - //! [maxBatchSize] or [numSequences] - TensorPtr eagleNetGenContextLengthsHost; - //! [maxBatchSize] or [numSequences] - TensorPtr eagleNetGenPastKeyValueLengthsHost; - //! [maxBatchSize * maxDecodingTokens] or [numSequences * maxDecodingTokens] - TensorPtr inputGenTokensHost; - //! [maxBatchSize] or [numSequences] - TensorPtr chunkedContextNextTokens; - //! [1] - TensorPtr useSpecDecoding; - - // For Eagle-2 - //! [1] - TensorPtr useDynamicTreeHost; - //! [1] - TensorPtr dynamicTreeMaxTopKHost; - //! [maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens] - TensorPtr prevScores; - //! [maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens] - TensorPtr currentExpandIndices; - //! [maxBatchSize, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] or [numSequences, - //! numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] - TensorPtr allLayersScores; - //! [maxBatchSize, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] or [numSequences, - //! numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] - TensorPtr allLayersDraftTokenIds; - //! [maxBatchSize, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] or [numSequences, - //! numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] - TensorPtr allLayersDraftTokenIdsPredecessor; - - void create(SizeType32 maxNumSequences, BufferManager const& manager, ModelConfig const& modelConfig, - WorldConfig const& worldConfig); - }; - - Inputs engineInputs; - - class EngineOutputs - { - public: - //! [batchSize, maxDecodingDraftTokens] - TensorPtr nextDraftTokens; - //! [batchSize] - TensorPtr nextDraftLens; - //! [batchSize, maxNumPaths, maxPathLen] - TensorPtr nextDraftPaths; - - //! [batchSize, maxPathLen] - TensorPtr acceptedTokens; - //! [batchSize] - TensorPtr acceptedLens; - //! [batchSize] - TensorPtr acceptedPaths; - //! [batchSize] - TensorPtr chunkedContextNextTokens; - - } engineOutputs; - -public: - EagleBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig); - - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, runtime::ModelConfig const& modelConfig); - - void setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::ITensor const& requestTypes, ITensor const& seqSlots, EagleBuffers::Inputs const& decoderBuffers, - runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) const; - - void insertInputTensors( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::WorldConfig const& worldConfig) const; - -private: - template - void setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, - SizeType32 vocabSizePadded, ITensor const& seqSlots, EagleBuffers::Inputs const& draftBuffers, - runtime::EagleModule const& eagleModule, runtime::BufferManager const& manager) const; - -private: - // helper tensors - std::size_t scanReduceTempStorageBytes{0}; - float mDefaultPosteriorThreshold{0.09f}; - bool mDoGreedySampling{true}; - BufferPtr scanReduceTempStorage; - TensorPtr cumSumGenerationLengths; - TensorPtr maxGenerationLength; - TensorPtr chunkedContextNextTokensHost; - TensorPtr greedySamplingHost; - TensorPtr posteriorAlphaHost; - TensorPtr posteriorThresholdHost; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/explicitDraftTokensBuffers.h b/cpp/include/tensorrt_llm/runtime/explicitDraftTokensBuffers.h deleted file mode 100644 index afb0f14a6d61..000000000000 --- a/cpp/include/tensorrt_llm/runtime/explicitDraftTokensBuffers.h +++ /dev/null @@ -1,146 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/explicitDraftTokensModule.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -namespace tensorrt_llm::runtime -{ - -class ExplicitDraftTokensBuffers -{ -public: - using SizeType32 = runtime::SizeType32; - using ITensor = runtime::ITensor; - using BufferPtr = runtime::IBuffer::SharedPtr; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - class Inputs - { - public: - //! [maxBatchSize] - TensorPtr temperatures; - //! [maxBatchSize] - TensorPtr positionIdsBase; - //! [maxBatchSize] or [numGenSequences] - TensorPtr generationLengths; - //! [maxBatchSize] - TensorPtr randomDataSample; - //! [maxBatchSize, maxNumPaths, maxPathDraftLen] or [numGenSequences, maxNumPaths, maxPathDraftLen] - TensorPtr randomDataValidation; - //! [maxBatchSize, maxNumPaths, maxPathLen] or [numGenSequences, maxNumPaths, maxPathLen] - TensorPtr draftTokens; - //! [maxBatchSize, maxNumPaths, maxPathLen] or [numGenSequences, maxNumPaths, maxPathLen] - TensorPtr draftIndices; - //! [maxBatchSize, maxNumPaths, maxPathDraftLen, vocabSize] - //! or [numGenSequences, maxNumPaths, maxPathDraftLen, vocabSize] - TensorPtr draftProbs; - //! [maxBatchSize, maxDecodingTokens, ceil(maxDecodingTokens / 32)] - //! or [numGenSequences, maxDecodingTokens, ceil(maxDecodingTokens / 32)] - TensorPtr packedMasks; - //! [maxBatchSize] or [numGenSequences] - TensorPtr positionIds; - // [1], on pinned - TensorPtr maxGenLengthHost; - // [maxBatchSize] - TensorPtr generationLengthsHost; - // [1], on cpu - TensorPtr useSpecDecoding; - - void create(SizeType32 maxNumSequences, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - }; - - class EngineInputs : public Inputs - { - public: - //! [numSequences], on gpu - TensorPtr requestTypesDevice; - //! [numGenSequences] - TensorPtr positionOffsets; - } engineInputs; - - class EngineOutputs - { - public: - //! [batchSize] - TensorPtr nextGenerationLengths; - //! [batchSize] - TensorPtr nextPositionOffsets; - //! [batchSize, maxDecodingTokens, maxDecodingTokens], bool - TensorPtr masks; - - //! [batchSize, maxNumPaths, maxPathLen] - TensorPtr nextDraftTokens; - //! [batchSize, maxNumPaths, maxPathLen] - TensorPtr nextDraftIndices; - //! [batchSize, maxNumPaths, maxDraftPathLen, vocabSize] - TensorPtr nextDraftProbs; - - //! [batchSize * maxDecodingTokens] - TensorPtr nextFlatTokens; - //! [batchSize] - TensorPtr bestPathLengths; - //! [batchSize] - TensorPtr bestPathIndices; - //! [1] - TensorPtr maxGenToken; - //! [1] - TensorPtr totalGenToken; - //! [batchSize * maxDecodingTokens] - TensorPtr packedPositionIds; - } engineOutputs; - -public: - ExplicitDraftTokensBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, runtime::ModelConfig const& modelConfig); - - void setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, runtime::ITensor const& requestTypes, - ITensor const& seqSlots, ExplicitDraftTokensBuffers::Inputs const& decoderBuffers, - ITensor const& contextPositionIds, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, runtime::BufferManager const& manager, - runtime::CudaStream const& stream) const; - - void insertInputTensors( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::WorldConfig const& worldConfig) const; - -private: - template - void setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 vocabSizePadded, - ITensor const& seqSlots, ExplicitDraftTokensBuffers::Inputs const& draftBuffers, - ITensor const& contextPositionIds, runtime::ExplicitDraftTokensModule const& explicitDraftTokensModule, - runtime::CudaStream const& stream) const; - -public: - // helper tensors - std::size_t scanTempStorageBytes{0}; - BufferPtr scanTempStorage; - TensorPtr cumSumGenerationLengths; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoder.h b/cpp/include/tensorrt_llm/runtime/gptDecoder.h deleted file mode 100644 index 5a785e84fe75..000000000000 --- a/cpp/include/tensorrt_llm/runtime/gptDecoder.h +++ /dev/null @@ -1,154 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/decodingInput.h" -#include "tensorrt_llm/runtime/decodingOutput.h" -#include "tensorrt_llm/runtime/samplingConfig.h" - -#include "tensorrt_llm/common/tllmDataType.h" -#include - -#include - -namespace tensorrt_llm -{ - -namespace layers -{ -// Forward declaration -template -class DynamicDecodeLayer; -} // namespace layers - -namespace runtime -{ - -class SpeculativeDecodingModule; - -class DecodingLayerWorkspace; - -class IGptDecoder -{ -public: - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorConstPtr = runtime::ITensor::SharedConstPtr; - - virtual ~IGptDecoder() = default; - - /// @param explicitDraftTokensDType is only used by ExplicitDraftTokens model to WAR the lack of bf16 decoder. - virtual void setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, - std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, - std::optional> const& lookaheadPrompt = std::nullopt, - std::optional> const& lookaheadAlgoConfigs = std::nullopt) - = 0; - - virtual void forwardAsync(DecodingOutput& output, DecodingInput const& input) = 0; - - virtual void forwardSync(DecodingOutput& output, DecodingInput const& input) = 0; - - virtual SamplingConfig const& getSamplingConfig() = 0; - - virtual void disableLookahead( - std::optional const& samplingConfig, SizeType32 batchSize, TensorConstPtr batchSlots) - = 0; - - static std::unique_ptr create(executor::DecodingMode const& mode, tensorrt_llm::DataType dtype, - size_t maxNumSequences, size_t maxBeamWidth, size_t vocabSize, size_t vocabSizePadded, - BufferManager::CudaStreamPtr const& stream, - std::shared_ptr const& speculativeDecodingModule = nullptr); -}; - -template -class GptDecoder : public virtual IGptDecoder -{ - -public: - using CudaStreamPtr = BufferManager::CudaStreamPtr; - using TensorPtr = std::shared_ptr; - - GptDecoder(executor::DecodingMode const& mode, size_t maxNumSequences, size_t maxBeamWidth, size_t vocabSize, - size_t vocabSizePadded, CudaStreamPtr const& stream, - std::shared_ptr speculativeDecodingModule = nullptr); - - void setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, - std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, - std::optional> const& lookaheadPrompt = std::nullopt, - std::optional> const& lookaheadAlgoConfigs - = std::nullopt) override; - - void forwardAsync(DecodingOutput& output, DecodingInput const& input) override; - - void forwardSync(DecodingOutput& output, DecodingInput const& input) override; - - SamplingConfig const& getSamplingConfig() override - { - return mSamplingConfig; - } - - void disableLookahead( - std::optional const& samplingConfig, SizeType32 batchSize, TensorConstPtr batchSlots) override; - -private: - std::shared_ptr mManager; - std::shared_ptr> mDynamicDecodeLayer; - std::shared_ptr mDecodingLayerWorkspace; - - SamplingConfig mSamplingConfig; - - size_t mMaxNumSequences; - size_t mVocabSize; - size_t mVocabSizePadded; - - executor::DecodingMode mDecodingMode; -}; - -inline std::unique_ptr IGptDecoder::create(executor::DecodingMode const& mode, - tensorrt_llm::DataType dtype, size_t maxNumSequences, size_t maxBeamWidth, size_t vocabSize, size_t vocabSizePadded, - BufferManager::CudaStreamPtr const& stream, - std::shared_ptr const& speculativeDecodingModule) -{ - switch (dtype) - { - case tensorrt_llm::DataType::kFLOAT: - return std::make_unique>( - mode, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, stream, speculativeDecodingModule); - case tensorrt_llm::DataType::kHALF: - return std::make_unique>( - mode, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, stream, speculativeDecodingModule); - default: - TLLM_THROW("Unsupported decoder data type: %d. Use either kFLOAT or kHALF.", static_cast(dtype)); - return nullptr; - } -} - -/// @brief Helper function to produce batch slots [0, 1, ..., batchSize - 1] for paths that do not explicitly provide -/// batch slots to the decoder. -inline runtime::ITensor::SharedConstPtr getDefaultBatchSlots(runtime::SizeType32 batchSize) -{ - auto defaultBatchSlots = runtime::BufferManager::pinnedPool( - runtime::ITensor::makeShape({batchSize}), runtime::TRTDataType::value); - auto range = runtime::BufferRange(*defaultBatchSlots); - std::iota(range.begin(), range.end(), 0); - return defaultBatchSlots; -} -} // namespace runtime -} // namespace tensorrt_llm diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h deleted file mode 100644 index d5447f441163..000000000000 --- a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h +++ /dev/null @@ -1,92 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaEvent.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/gptDecoder.h" -#include "tensorrt_llm/runtime/iGptDecoderBatched.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include - -namespace tensorrt_llm::batch_manager -{ -class LlmRequest; -} // namespace tensorrt_llm::batch_manager - -namespace tensorrt_llm::runtime -{ - -//! GPT decoder class with support for in-flight batching -class GptDecoderBatched : public IGptDecoderBatched -{ -public: - using CudaStreamPtr = std::shared_ptr; - using LlmRequestPtr = std::shared_ptr; - using RequestVector = std::vector; - using TensorPtr = ITensor::SharedPtr; - - explicit GptDecoderBatched(CudaStreamPtr stream); - - void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) override; - - void disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) override; - - CudaEvent forwardAsync( - decoder::DecoderState const& decoderState, batch_manager::DecoderInputBuffers const& input) override; - void forward(decoder::DecoderState const& decoderState, batch_manager::DecoderInputBuffers const& input) override; - - //! @brief Gather final beam search results for request `batchSlot`. - //! Result will only be available after event returned. - [[nodiscard]] CudaEvent finalize(decoder::DecoderState const& decoderState, SizeType32 batchSlot, - SamplingConfig const& samplingConfig, bool streaming) const override; - - CudaStreamPtr getDecoderStream() const - { - return mDecoderStream; - } - - IGptDecoder& getUnderlyingDecoder() const - { - return *mDecoder.get(); - } - - [[nodiscard]] BufferManager const& getBufferManager() const - { - return mBufferManager; - } - -private: - //! @brief Calls decoders for tokens per engine step - void forwardDispatch(decoder::DecoderState const& decoderState, batch_manager::DecoderInputBuffers const& input); - -private: - CudaStreamPtr mRuntimeStream; - CudaStreamPtr mDecoderStream; - BufferManager mBufferManager; - - using GptDecoderPtr = std::unique_ptr; - GptDecoderPtr mDecoder; -}; -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h deleted file mode 100644 index b664bc007f0e..000000000000 --- a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h +++ /dev/null @@ -1,81 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/runtime/cudaEvent.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include - -namespace tensorrt_llm::batch_manager -{ -class DecoderInputBuffers; -class LlmRequest; -} // namespace tensorrt_llm::batch_manager - -namespace tensorrt_llm::runtime -{ -class SamplingConfig; - -namespace decoder -{ -class DecoderState; -} - -//! GPT decoder class with support for in-flight batching -class IGptDecoderBatched -{ -public: - using CudaStreamPtr = std::shared_ptr; - using LlmRequestPtr = std::shared_ptr; - using RequestVector = std::vector; - using TensorPtr = std::shared_ptr; - - //! @brief Setup the decoder before calling `forward()` - virtual void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) - = 0; - - //! @brief Disable Lookahead decoding. - virtual void disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) = 0; - - //! @brief Run one step for all requests without blocking the host process and return the token for synchronization. - virtual CudaEvent forwardAsync( - decoder::DecoderState const& decoderState, batch_manager::DecoderInputBuffers const& input) - = 0; - - //! @brief Run one step for all requests and wait for completion on the host. - virtual void forward(decoder::DecoderState const& decoderState, batch_manager::DecoderInputBuffers const& input) - = 0; - - //! @brief Gather final beam search results for request `batchIdx`. - //! Result will only be available after event returned - [[nodiscard]] virtual CudaEvent finalize(decoder::DecoderState const& decoderState, SizeType32 batchSlot, - SamplingConfig const& samplingConfig, bool streaming) const - = 0; - -protected: - IGptDecoderBatched() = default; - virtual ~IGptDecoderBatched() = default; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/samplingConfig.h b/cpp/include/tensorrt_llm/runtime/samplingConfig.h index 03355167653f..1945cafd84a3 100644 --- a/cpp/include/tensorrt_llm/runtime/samplingConfig.h +++ b/cpp/include/tensorrt_llm/runtime/samplingConfig.h @@ -368,7 +368,7 @@ class SamplingConfig OptVec earlyStopping; // [1] or [batchSize] OptVec> beamWidthArray; // [maxBeamWidthArrayLength] or [batchSize, maxBeamWidthArrayLength] - // speculative decoding, only the first value is used (in gptDecoderBatched.cpp) + // speculative decoding OptVec draftAcceptanceThreshold; // [1] or [batchSize] // medusa params diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 691c49e545aa..692b8c082e44 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -142,7 +142,6 @@ endif() add_subdirectory(common) add_subdirectory(kernels) -add_subdirectory(layers) add_subdirectory(runtime) set(BATCH_MANAGER_TARGET tensorrt_llm_batch_manager_static) @@ -197,7 +196,6 @@ set(TRTLLM_LINK_LIBS gemm_swiglu_sm90_src cutlass_src cute_dsl_src - layers_src runtime_src compressorKernels_src mhcKernels_src diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index f61e58e16b28..6062f10e80c2 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -28,18 +28,15 @@ set(SRCS mlaCacheFormatter.cpp cacheTransceiver.cpp capacityScheduler.cpp - createNewDecoderRequests.cpp contextProgress.cpp contextTransferCoordinator.cpp dataTransceiver.cpp - decoderBuffers.cpp kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp kvCacheManagerV2Utils.cu llmRequest.cpp - medusaBuffers.cpp microBatchScheduler.cpp pauseRequests.cpp peftCacheManager.cpp diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp deleted file mode 100644 index 2d090a5612a5..000000000000 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ /dev/null @@ -1,750 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/medusaBuffers.h" -#include "tensorrt_llm/batch_manager/utils/logitsThread.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/decodingInput.h" -#include "tensorrt_llm/runtime/decodingOutput.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/speculativeDecodingMode.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h" - -#include "tensorrt_llm/common/tllmDataType.h" - -using namespace tensorrt_llm::runtime; - -namespace tc = tensorrt_llm::common; -namespace te = tensorrt_llm::executor; -namespace tr = tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -using SizeType32 = CreateNewDecoderRequests::SizeType32; -using TensorPtr = CreateNewDecoderRequests::TensorPtr; -using SharedConstPtr = CreateNewDecoderRequests::SharedConstPtr; -template -using OptionalRef = tensorrt_llm::common::OptionalRef; - -namespace -{ - -void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffers& inputBuffers, - 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); - - auto batchSlotsRange = tr::BufferRange(*batchSlotsView); - auto fillValuesRange = tr::BufferRange(*fillValuesView); - - // fill buffers on host - SizeType32 batchIdx{0}; - for (auto const& llmReq : contextRequests) - { - auto const currentSequenceLen - = llmReq->mPromptLen + llmReq->getMaxNumGeneratedTokens() + llmReq->getNumContextPhaseGenerationTokens(); - // Get position of the current sequence in the decoder - auto const seqSlot = llmReq->mSeqSlot.value(); - batchSlotsRange[batchIdx] = seqSlot; - fillValuesRange[batchIdx] = currentSequenceLen; - ++batchIdx; - } - - // copy sequence lengths - { - auto batchSlotsDeviceView = tr::ITensor::slice(inputBuffers.setupBatchSlotsDevice, 0, batchSize); - auto fillValuesViewDevice = tr::ITensor::slice(inputBuffers.fillValuesDevice, 0, batchSize); - - bufferManager.copy(*batchSlotsView, *batchSlotsDeviceView); - bufferManager.copy(*fillValuesView, *fillValuesViewDevice); - tr::kernels::invokeFillBatch(sequenceLengths, *batchSlotsDeviceView, beamWidth, *fillValuesViewDevice, stream); - } -} - -/// @brief Retrieve the embedding bias from the request. This potentially makes a copy of the tensor -/// to the appropriate type if the input tensor does not match it. -[[nodiscard]] TensorPtr getEmbeddingBias(tensorrt_llm::DataType logitsType, TensorPtr const& tensor) -{ - // Check that embedding bias type is same as logits type. If so, we can return the tensor right away - if (tensor->getDataType() == logitsType) - { - return tensor; - } - - // Support FP32 input for FP16 embedding bias (in the case of FP8 models) - if (tensor->getDataType() == tensorrt_llm::DataType::kFLOAT && logitsType == tensorrt_llm::DataType::kHALF) - { - // Do a deep copy of the tensor to the expected type - TLLM_LOG_WARNING( - "Embedding bias data type must be same as model logits type, will copy the tensor from float to half"); - - TLLM_CHECK_WITH_INFO( - tensor->getMemoryType() != MemoryType::kGPU, "Embedding bias tensor needs to be in CPU memory for casting"); - - auto const shape = tensor->getShape(); - TLLM_CHECK(shape.nbDims == 2); // [1, vocabSizePadded] - TLLM_CHECK(shape.d[0] == 1); - auto newTensor = tensorrt_llm::runtime::BufferManager::pinnedPool(shape, logitsType); - - auto const tensorRange = BufferRange(*tensor); - auto newTensorRange = BufferRange(*newTensor); - - std::transform(tensorRange.begin(), tensorRange.end(), newTensorRange.begin(), - [](float value) -> half { return static_cast(value); }); - - return newTensor; - } - - TLLM_THROW("Embedding bias data type must be same as model logits type."); -} - -} // namespace - -std::tuple, std::vector, - std::vector> -CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - tensorrt_llm::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); - - RequestVector finishedContextRequests; - std::copy_if(contextRequests.begin(), contextRequests.end(), std::back_inserter(finishedContextRequests), - [](auto const& llmReq) { return llmReq->isLastContextChunk(); }); - - if (!finishedContextRequests.empty()) - { - copySequenceLengths( - finishedContextRequests, inputBuffers, *decoderState.getSequenceLengths(), beamWidth, runtimeStream); - } - - auto [lookaheadPrompt, lookaheadAlgoConfigs] - = createDecoderRequests(finishedContextRequests, inputBuffers.inputsIds, decodingConfig, decoderState, - logitsType, modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, medusaBuffers); - - auto const batchSize = finishedContextRequests.size(); - - std::vector samplingConfigs; - samplingConfigs.reserve(batchSize); - for (auto const& llmReq : finishedContextRequests) - { - samplingConfigs.push_back(llmReq->mSamplingConfig); - } - - TensorPtr batchSlotsView = runtime::ITensor::slice(inputBuffers.setupBatchSlots, 0, batchSize); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {std::move(batchSlotsView), std::move(samplingConfigs), std::move(lookaheadPrompt), - std::move(lookaheadAlgoConfigs)}; -} - -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) -{ - 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, tensorrt_llm::DataType logitsType, - runtime::ModelConfig const& modelConfig, BufferManager const& manager) -{ - TensorPtr const embeddingBiasSlice = ITensor::slice(constPointerCast(dJointInput.embeddingBias), batchSlot, 1); - 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(), - "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); - } -} - -void setupWords(std::vector& jointWordsLists, - std::optional const& requestWordsList, SharedConstPtr& jointWordsPtrs, SharedConstPtr& jointWordsLens, - SizeType32& jointMaxWordsLen, SizeType32 batchSlot, BufferManager const& manager) -{ - if (requestWordsList.has_value()) - { - // 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 - { - runtime::bufferCast(*constPointerCast(jointWordsLens))[batchSlot] = 0; - } -}; - -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); - } -} - -void initializeOutputs(DecodingOutput& dJointOutput, SizeType32 batchSlot, SizeType32 maxDecodingEngineTokens, - BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto finishedSum = ITensor::slice(dJointOutput.finishedSum, batchSlot, 1); - manager.setZero(*finishedSum); - - for (SizeType32 ti = 0; ti < maxDecodingEngineTokens; ++ti) - { - TensorPtr const newTokensStepView = ITensor::slice(dJointOutput.newTokensSteps, ti, 1); - newTokensStepView->squeeze(0); - auto newTokensVec = ITensor::slice(newTokensStepView, batchSlot, 1); - manager.setZero(*newTokensVec); - } - - TensorPtr const finishedStepsSlice = ITensor::slice(dJointOutput.finishReasons, batchSlot, 1); - manager.setZero(*finishedStepsSlice); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void retrieveDraftLogits(TensorPtr& draftLogitsHost, std::shared_ptr const& reqDraftLogits, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, bool speculativeDecodingFastLogits, - bool isLeaderInOrchMode, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (!speculativeDecodingFastLogits) - { - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - bufferManager.copy(*reqDraftLogits, *draftLogitsHost); - return; - } - - if (isLeaderInOrchMode) - { - // reqDraftLogits contains metadata for fast-logits path; validate size. - auto constexpr fastLogitsInfoSize = sizeof(te::SpeculativeDecodingFastLogitsInfo); - TLLM_CHECK_WITH_INFO(reqDraftLogits->getSizeInBytes() >= fastLogitsInfoSize, - "Draft logits metadata buffer is too small to hold SpeculativeDecodingFastLogitsInfo."); - te::SpeculativeDecodingFastLogitsInfo fastLogitsInfo{}; - std::memcpy(&fastLogitsInfo, reqDraftLogits->data(), fastLogitsInfoSize); - utils::targetModelReceiveLogits(draftLogitsHost, fastLogitsInfo, modelConfig.getLogitsDtype()); - - // Broadcast to other ranks if needed - if (worldConfig.isTensorParallel()) - { - auto const& commSession = COMM_SESSION; - auto shape = draftLogitsHost->getShape(); - commSession.bcastValue(shape.d[0], 0); - commSession.bcastValue(shape.d[1], 0); - commSession.bcast(draftLogitsHost->data(), draftLogitsHost->getSizeInBytes(), mpi::MpiType::kUINT8, 0); - } - } - else - { - TLLM_CHECK_WITH_INFO(worldConfig.isTensorParallel(), - "Fast logits path requires tensor-parallel broadcast for non-leader ranks."); - - // Get logits from leader rank - auto const& commSession = COMM_SESSION; - int64_t dims[2]; - commSession.bcastValue(dims[0], 0); - commSession.bcastValue(dims[1], 0); - draftLogitsHost->reshape(ITensor::makeShape({dims[0], dims[1]})); - commSession.bcast(draftLogitsHost->data(), draftLogitsHost->getSizeInBytes(), mpi::MpiType::kUINT8, 0); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -}; - -//! @brief Setups decoder internal tensors for new request in Draft model Sps mode -void newRequestDraftTokensExternal(DecodingInput& jointDecodingInput, SizeType32 batchIdx, LlmRequest const& llmReq, - SizeType32 numDecodingEngineTokens, runtime::ModelConfig const& modelConfig, WorldConfig const& worldConfig, - bool speculativeDecodingFastLogits, bool isLeaderInOrchMode, CudaStream const& decoderStream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - BufferManager decoderBufferManager{std::make_shared(decoderStream.get())}; - - TLLM_CHECK(jointDecodingInput.externalDraftTokensInputs); - auto& externalDraftTokensInputs = jointDecodingInput.externalDraftTokensInputs; - - auto const& draftTokens = llmReq.getDraftTokens(); - auto const numDraftTokens = numDecodingEngineTokens - 1; - - auto numDraftTokensHostRange = runtime::BufferRange(*externalDraftTokensInputs->numDraftTokensHost); - numDraftTokensHostRange[batchIdx] = numDraftTokens; - auto numDraftTokensView = ITensor::slice(externalDraftTokensInputs->numDraftTokens, batchIdx, 1); - runtime::kernels::invokeFill(*numDraftTokensView, numDraftTokens, decoderStream); - - if (numDraftTokens > 0) - { - TensorPtr draftTokenIdsHostSlice - = ITensor::slice(externalDraftTokensInputs->draftTokenIdsHost, {batchIdx, 0}, numDraftTokens); - // Copy to pinned host memory (don't care about stream of bufferManager) - decoderBufferManager.copy(draftTokens->data(), *draftTokenIdsHostSlice); - - TensorPtr draftTokenIdsSlice - = ITensor::slice(externalDraftTokensInputs->draftTokenIds, {batchIdx, 0}, numDraftTokens); - decoderBufferManager.copy(*draftTokenIdsHostSlice, *draftTokenIdsSlice); - } - - auto const& draftLogits = llmReq.getDraftLogits(); - auto const useDraftLogits = draftLogits.has_value(); - - auto useDraftLogitsHostRange = runtime::BufferRange(*externalDraftTokensInputs->useDraftLogitsHost); - useDraftLogitsHostRange[batchIdx] = useDraftLogits; - auto useDraftLogitsView = ITensor::slice(externalDraftTokensInputs->useDraftLogits, batchIdx, 1); - runtime::kernels::invokeFill(*useDraftLogitsView, useDraftLogits, decoderStream); - - if (useDraftLogits) - { - TensorPtr draftLogitsHostSlice - = ITensor::slice(externalDraftTokensInputs->draftLogitsHost, {batchIdx, 0}, numDraftTokens); - retrieveDraftLogits(draftLogitsHostSlice, draftLogits.value(), modelConfig, worldConfig, - speculativeDecodingFastLogits, isLeaderInOrchMode, decoderBufferManager); - - TensorPtr draftLogitsSlice - = ITensor::slice(externalDraftTokensInputs->draftLogits, {batchIdx, 0}, numDraftTokens); - decoderBufferManager.copy(*draftLogitsHostSlice, *draftLogitsSlice); - } - - auto const& samplingConfig = llmReq.mSamplingConfig; - bool const useRandomAcceptanceThreshold = !samplingConfig.draftAcceptanceThreshold.has_value(); - float const constantThreshold - = useRandomAcceptanceThreshold ? 0 : samplingConfig.draftAcceptanceThreshold.value()[0]; - - externalDraftTokensInputs->useRandomAcceptanceThreshold = useRandomAcceptanceThreshold; - externalDraftTokensInputs->constantThreshold = constantThreshold; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Setups decoder internal tensors for new Medusa request -void newRequestMedusa(DecodingInput& jointDecodingInput, SizeType32 batchIdx, LlmRequest& llmReq, - SizeType32 numDecodingEngineTokens, SizeType32 maxDecodingEngineTokens, MedusaBuffers const& medusaBuffers, - CudaStream const& decoderStream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - 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. - auto medusaPaths = ITensor::slice(medusaBuffers.medusaPathsDevice, 0, 1); - auto medusaTreeIds = ITensor::slice(medusaBuffers.medusaTreeIdsDevice, 0, 1); - - BufferManager manager{std::make_shared(decoderStream.get())}; - - auto& medusaInputs = jointDecodingInput.medusaInputs; - - TensorPtr curTokensPerStepSlice - = ITensor::slice(constPointerCast(medusaInputs->medusaCurTokensPerStep), batchIdx, 1); - // Context phase Medusa processes 1 token only, new value from targetTokensPerStep will be filled at the end - // of first decoder - runtime::kernels::invokeFill(*curTokensPerStepSlice, 1, decoderStream); - TensorPtr targetTokensPerStepSlice - = ITensor::slice(constPointerCast(medusaInputs->medusaTargetTokensPerStep), batchIdx, 1); - TLLM_CHECK_WITH_INFO(numDecodingEngineTokens <= maxDecodingEngineTokens, - "Tokens per step for (%d) is larger than maximum tokens per step (%d)", numDecodingEngineTokens, - maxDecodingEngineTokens); - runtime::kernels::invokeFill(*targetTokensPerStepSlice, numDecodingEngineTokens, decoderStream); - - TensorPtr pathsSlice = ITensor::slice(constPointerCast(medusaInputs->medusaPaths), batchIdx, 1); - manager.copy(*medusaPaths, *pathsSlice); - - TensorPtr treeIdsSlice = ITensor::slice(constPointerCast(medusaInputs->medusaTreeIds), batchIdx, 1); - manager.copy(*medusaTreeIds, *treeIdsSlice); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Setups decoder internal tensors for new Lookahead request -void newRequestLookahead(DecodingInput& jointDecodingInput, DecodingOutput& jointDecodingOutput, SizeType32 batchIdx, - CudaStream const& runtimeStream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(jointDecodingOutput.lookaheadOutputs); - TLLM_CHECK(jointDecodingInput.lookaheadInputs); - - // The first generation step only generate 1 token. - TensorPtr curTokensPerStepSlice - = ITensor::slice(constPointerCast(jointDecodingInput.lookaheadInputs->tokensPerStep), batchIdx, 1); - runtime::kernels::invokeFill(*curTokensPerStepSlice, 1, runtimeStream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Setups decoder internal tensors for new Explicit draft tokens request -void newRequestExplicitDraftTokens( - DecodingOutput& jointDecodingOutput, SizeType32 batchIdx, LlmRequest const& llmReq, CudaStream const& runtimeStream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(jointDecodingOutput.explicitDraftTokensBuffers); - - auto const inputLen = llmReq.getPromptLen(); - - TensorPtr positionIdsBaseSlice - = ITensor::slice(jointDecodingOutput.explicitDraftTokensBuffers->positionIdsBase, batchIdx, 1); - runtime::kernels::invokeFill(*positionIdsBaseSlice, inputLen, runtimeStream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Setups decoder internal tensors for new Eagle request -void newRequestEagle(DecodingOutput& jointDecodingOutput, SizeType32 batchIdx, LlmRequest const& llmReq, - runtime::ModelConfig const& modelConfig, executor::DecodingConfig const& decodingConfig, - CudaStream const& runtimeStream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(jointDecodingOutput.eagleBuffers); - auto& eagleBuffers = *jointDecodingOutput.eagleBuffers; - - auto const inputLen = llmReq.getPromptLen(); - - BufferManager manager{std::make_shared(runtimeStream.get())}; - - TensorPtr eagleNetCtxRequestTypesHostSlice = ITensor::slice(eagleBuffers.eagleNetCtxRequestTypesHost, batchIdx, 1); - TensorPtr eagleNetCtxContextLengthsHostSlice - = ITensor::slice(eagleBuffers.eagleNetCtxContextLengthsHost, batchIdx, 1); - TensorPtr eagleNetCtxPastKeyValueLengthsHostSlice - = ITensor::slice(eagleBuffers.eagleNetCtxPastKeyValueLengthsHost, batchIdx, 1); - - runtime::bufferCast(*eagleNetCtxRequestTypesHostSlice)[0] = 0; - runtime::bufferCast(*eagleNetCtxContextLengthsHostSlice)[0] = inputLen; - runtime::bufferCast(*eagleNetCtxPastKeyValueLengthsHostSlice)[0] = inputLen; - - TensorPtr eagleNetGenRequestTypesHostSlice = ITensor::slice(eagleBuffers.eagleNetGenRequestTypesHost, batchIdx, 1); - TensorPtr eagleNetGenContextLengthsHostSlice - = ITensor::slice(eagleBuffers.eagleNetGenContextLengthsHost, batchIdx, 1); - TensorPtr eagleNetGenPastKeyValueLengthsHostSlice - = ITensor::slice(eagleBuffers.eagleNetGenPastKeyValueLengthsHost, batchIdx, 1); - - runtime::bufferCast(*eagleNetGenRequestTypesHostSlice)[0] = 1; - runtime::bufferCast(*eagleNetGenContextLengthsHostSlice)[0] = inputLen; - runtime::bufferCast(*eagleNetGenPastKeyValueLengthsHostSlice)[0] = inputLen; - - auto const eagleModule = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - std::optional eagleChoicesOpt; - - auto const& eagleConfig = llmReq.getEagleConfig() ? llmReq.getEagleConfig() : decodingConfig.getEagleConfig(); - - if (eagleConfig) - { - eagleChoicesOpt = eagleConfig->getEagleChoices(); - } - - if (!eagleConfig || !eagleConfig->useDynamicTree()) - { - TensorPtr draftPathsHostSlice = ITensor::slice(eagleBuffers.draftPathsHost, batchIdx, 1); - TensorPtr draftPathsSlice = ITensor::slice(eagleBuffers.draftPaths, batchIdx, 1); - - // eagleConfig is nullptr or Eagle-1 - std::vector topKs; - auto const depth = runtime::utils::initTensorsFromChoices(modelConfig.getSpeculativeDecodingModule(), - eagleChoicesOpt.value_or(eagleModule->getDefaultEagleChoices()), topKs, nullptr, nullptr, nullptr, - draftPathsHostSlice, nullptr, {eagleModule->getMaxNonLeafNodesPerLayer()}); - TLLM_CHECK_WITH_INFO(depth == modelConfig.getSpeculativeDecodingModule().getMaxDraftPathLen(), - "EAGLE-1 requires Eagle-tree depth being equal to the the number of build-time EAGLE layers."); - - manager.copy(*draftPathsHostSlice, *draftPathsSlice); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Setups decoder internal tensors for new speculative decoding request -void newRequestSpeculativeDecoding(DecodingInput& jointDecodingInput, DecodingOutput& jointDecodingOutput, - SizeType32 batchIdx, LlmRequest& llmReq, SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 numDecodingEngineTokens, SizeType32 maxDecodingEngineTokens, - OptionalRef medusaBuffers, runtime::ModelConfig const& modelConfig, - WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, bool speculativeDecodingFastLogits, - bool isLeaderInOrchMode, CudaStream const& runtimeStream, CudaStream const& decoderStream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (speculativeDecodingMode.predictsDraftTokens()) - { - BufferManager manager{std::make_shared(decoderStream.get())}; - - TLLM_CHECK(jointDecodingOutput.speculativeDecodingOutputs); - auto& speculativeDecodingOutputs = *jointDecodingOutput.speculativeDecodingOutputs; - - TensorPtr nextDraftTokens = ITensor::slice(speculativeDecodingOutputs.nextDraftTokens, batchIdx, 1); - // FIXME: can we skip this? - manager.setZero(*nextDraftTokens); - if (speculativeDecodingMode.variableDraftLength()) - { - TensorPtr nextDraftTokensLen = ITensor::slice(speculativeDecodingOutputs.nextDraftTokensLen, batchIdx, 1); - manager.setZero(*nextDraftTokensLen); - } - } - - if (speculativeDecodingMode.isDraftTokensExternal()) - { - newRequestDraftTokensExternal(jointDecodingInput, batchIdx, llmReq, numDecodingEngineTokens, modelConfig, - worldConfig, speculativeDecodingFastLogits, isLeaderInOrchMode, decoderStream); - } - else if (speculativeDecodingMode.isMedusa()) - { - TLLM_CHECK(medusaBuffers); - newRequestMedusa(jointDecodingInput, batchIdx, llmReq, numDecodingEngineTokens, maxDecodingEngineTokens, - medusaBuffers.value(), decoderStream); - } - else if (speculativeDecodingMode.isLookaheadDecoding()) - { - newRequestLookahead(jointDecodingInput, jointDecodingOutput, batchIdx, runtimeStream); - } - else if (speculativeDecodingMode.isExplicitDraftTokens()) - { - newRequestExplicitDraftTokens(jointDecodingOutput, batchIdx, llmReq, runtimeStream); - } - else if (speculativeDecodingMode.isEagle()) - { - newRequestEagle(jointDecodingOutput, batchIdx, llmReq, modelConfig, decodingConfig, runtimeStream); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace - -std::tuple, std::vector> -CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedContextRequests, TensorPtr const& inputIds, - executor::DecodingConfig const& decodingConfig, runtime::decoder::DecoderState& decoderState, - tensorrt_llm::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) - { - auto const& reqTokens = llmReq->getTokens(0); - decoderInputSize += reqTokens.size(); - } - inputIds->resize(decoderInputSize); - - std::vector lookaheadPrompt; - std::vector lookaheadAlgoConfigs; - if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - TLLM_CHECK_WITH_INFO( - decodingConfig.getLookaheadDecodingConfig().has_value(), "Lookahead decoding config must be provided"); - lookaheadPrompt.reserve(finishedContextRequests.size()); - lookaheadAlgoConfigs.reserve(finishedContextRequests.size()); - } - - 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(); - - SizeType32 numDecodingEngineTokens{1}; - if (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) - { - numDecodingEngineTokens = llmReq->getNumDraftTokens() + 1; - } - else if (!modelConfig.getSpeculativeDecodingMode().isNone()) - { - numDecodingEngineTokens = modelConfig.getMaxDecodingTokens(); - } - - auto& dJointInput = decoderState.getJointDecodingInput(); - - 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) - { - initializeBeamSearch( - dJointInput, dJointOutput, batchSlot, endId, beamWidth, maxSequenceLength, decoderBufferManager); - } - - // Speculative execution - if (!decoderState.getSpeculativeDecodingMode().isNone()) - { - TLLM_CHECK(beamWidth == 1); - - if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - lookaheadPrompt.emplace_back(requestIds); - - auto const& lookaheadRuntimeConfig - = llmReq->getLookaheadConfig().value_or(decodingConfig.getLookaheadDecodingConfig().value()); - lookaheadAlgoConfigs.emplace_back(lookaheadRuntimeConfig); - } - - newRequestSpeculativeDecoding(decoderState.getJointDecodingInput(), decoderState.getJointDecodingOutput(), - batchSlot, *llmReq, decoderState.getSpeculativeDecodingMode(), numDecodingEngineTokens, - decoderState.getMaxDecodingEngineTokens(), medusaBuffers, modelConfig, worldConfig, decodingConfig, - mSpeculativeDecodingFastLogits, mIsLeaderInOrchMode, runtimeStream, decoderStream); - } - - inputOffset += promptLen; - } - - return {std::move(lookaheadPrompt), std::move(lookaheadAlgoConfigs)}; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp deleted file mode 100644 index fecc0851d361..000000000000 --- a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp +++ /dev/null @@ -1,325 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/batch_manager/decoderBuffers.h" - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/utils/mpiTags.h" - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -DecoderInputBuffers::DecoderInputBuffers( - SizeType32 maxBatchSize, SizeType32 maxDecoderSteps, BufferManager const& manager) -{ - auto const maxBatchSizeShape = ITensor::makeShape({maxBatchSize}); - auto const nvSizeType = TRTDataType::value; - - inputsIds = BufferManager::pinnedPool(ITensor::makeShape({0}), TRTDataType::value); - - setupBatchSlots = BufferManager::pinnedPool(maxBatchSizeShape, nvSizeType); - setupBatchSlotsDevice = manager.gpu(maxBatchSizeShape, nvSizeType); - - fillValues = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvSizeType); - fillValuesDevice = manager.gpu(maxBatchSizeShape, nvSizeType); - - forwardBatchSlots.reserve(maxDecoderSteps); - for (SizeType32 i = 0; i < maxDecoderSteps; ++i) - { - forwardBatchSlots.emplace_back(BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize}), nvSizeType)); - } -} - -void DecoderInputBuffers::setupMedusaLogits(SizeType32 maxNumSequences, ModelConfig const& modelConfig) -{ - if (modelConfig.getSpeculativeDecodingMode().isMedusa()) - { - auto const maxDraftPathLen = modelConfig.getSpeculativeDecodingModule().getMaxDraftPathLen(); - predictedDraftLogits.resize(maxNumSequences); - for (auto& medusaLogitsHead : predictedDraftLogits) - { - medusaLogitsHead.resize(maxDraftPathLen); - } - } -} - -DecoderOutputBuffers::DecoderOutputBuffers(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxSeqLen, - SizeType32 maxTokensPerStep, BufferManager const& manager) -{ - auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - - sequenceLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kINT32); - - finishedSumHost = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - - newOutputTokensHost - = BufferManager::pinned(ITensor::makeShape({maxTokensPerStep, maxNumSequences, maxBeamWidth}), TRTTokenIdType); - - cumLogProbsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); - - logProbsHost = BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); - - finishReasonsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kUINT8); -} - -void DecoderOutputBuffers::enableLookaheadDecoding(SizeType32 maxNumSequences, SizeType32 maxTokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - newOutputTokensHost->reshape(ITensor::makeShape({maxTokensPerStep, maxNumSequences, 1})); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderOutputBuffers::disableLookaheadDecoding(SizeType32 maxNumSequences) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - newOutputTokensHost->reshape(ITensor::makeShape({1, maxNumSequences, 1})); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderOutputBuffers::setupSpeculativeDecoding( - SizeType32 maxNumSequences, SizeType32 maxTokensPerStep, ModelConfig const& modelConfig) -{ - auto const speculativeDecodingMode = modelConfig.getSpeculativeDecodingMode(); - - auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - - if (speculativeDecodingMode.predictsDraftTokens()) - { - nextDraftTokensHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxTokensPerStep - 1}), TRTTokenIdType); - if (speculativeDecodingMode.variableDraftLength()) - { - nextDraftTokensLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - prevDraftTokensLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - } - } -} - -DecoderStepAsyncSend::DecoderStepAsyncSend(DecoderOutputBuffers const& decoderOutputBuffers, - runtime::decoder::DecoderState const& decoderState, bool const returnLogProbs, SizeType32 const maxBeamWidth, - bool const useMedusa, mpi::MpiComm const& commSession, int peer) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start send outputs of decoder to rank %d", peer); - - mRequest1 = commSession.sendAsync( - *decoderOutputBuffers.newOutputTokensHost, peer, mpi::MpiTag::kDecoderStepNewOutputTokensHost); - mRequest2 - = commSession.sendAsync(*decoderOutputBuffers.finishedSumHost, peer, mpi::MpiTag::kDecoderStepFinishedSumHost); - mRequest3 = commSession.sendAsync( - *decoderOutputBuffers.sequenceLengthsHost, peer, mpi::MpiTag::kDecoderStepSequenceLengthsHost); - mRequest4 = returnLogProbs - ? commSession.sendAsync(*decoderOutputBuffers.cumLogProbsHost, peer, mpi::MpiTag::kDecoderStepCumLogProbsHost) - : nullptr; - mRequest5 = returnLogProbs - ? commSession.sendAsync(*decoderOutputBuffers.logProbsHost, peer, mpi::MpiTag::kDecoderStepLogProbsHost) - : nullptr; - mRequest6 = maxBeamWidth > 1 ? commSession.sendAsync( - *decoderState.getCacheIndirectionOutput(), peer, mpi::MpiTag::kDecoderStepCacheIndirectionOutput) - : nullptr; - mRequest7 = useMedusa ? commSession.sendAsync(*decoderState.getAcceptedLengthsCumSum(), peer, - mpi::MpiTag::kDecoderStepAcceptedLengthsCumSumDevice) - : nullptr; - mRequest8 = useMedusa ? commSession.sendAsync( - *decoderState.getAcceptedPackedPaths(), peer, mpi::MpiTag::kDecoderStepAcceptedPackedPathsDevice) - : nullptr; - mRequest9 = commSession.sendAsync( - *decoderOutputBuffers.finishReasonsHost, peer, mpi::MpiTag::kDecoderStepFinishReasonsHost); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderStepAsyncSend::recv(DecoderOutputBuffers const& decoderOutputBuffers, - runtime::decoder::DecoderState const& decoderState, bool const returnLogProbs, SizeType32 const maxBeamWidth, - bool const useMedusa, mpi::MpiComm const& commSession, int const peer) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start recv outputs of decoder from rank %d", peer); - - commSession.recv(*decoderOutputBuffers.newOutputTokensHost, peer, mpi::MpiTag::kDecoderStepNewOutputTokensHost); - commSession.recv(*decoderOutputBuffers.finishedSumHost, peer, mpi::MpiTag::kDecoderStepFinishedSumHost); - commSession.recv(*decoderOutputBuffers.sequenceLengthsHost, peer, mpi::MpiTag::kDecoderStepSequenceLengthsHost); - if (returnLogProbs) - { - commSession.recv(*decoderOutputBuffers.cumLogProbsHost, peer, mpi::MpiTag::kDecoderStepCumLogProbsHost); - commSession.recv(*decoderOutputBuffers.logProbsHost, peer, mpi::MpiTag::kDecoderStepLogProbsHost); - } - if (maxBeamWidth > 1) - { - commSession.recv( - *decoderState.getCacheIndirectionOutput(), peer, mpi::MpiTag::kDecoderStepCacheIndirectionOutput); - } - if (useMedusa) - { - commSession.recv( - *decoderState.getAcceptedLengthsCumSum(), peer, mpi::MpiTag::kDecoderStepAcceptedLengthsCumSumDevice); - commSession.recv( - *decoderState.getAcceptedPackedPaths(), peer, mpi::MpiTag::kDecoderStepAcceptedPackedPathsDevice); - } - commSession.recv(*decoderOutputBuffers.finishReasonsHost, peer, mpi::MpiTag::kDecoderStepFinishReasonsHost); - - TLLM_LOG_DEBUG("end recv outputs of decoder from rank %d", peer); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -DecoderStepAsyncSend::~DecoderStepAsyncSend() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mRequest1->wait(); - mRequest2->wait(); - mRequest3->wait(); - if (mRequest4) - mRequest4->wait(); - if (mRequest5) - mRequest5->wait(); - if (mRequest6) - mRequest6->wait(); - if (mRequest7) - mRequest7->wait(); - if (mRequest8) - mRequest8->wait(); - mRequest9->wait(); - - TLLM_LOG_DEBUG("end send outputs of decoder"); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderStepAsyncSend::bcast(DecoderOutputBuffers const& decoderOutputBuffers, - runtime::decoder::DecoderState const& decoderState, bool const returnLogProbs, SizeType32 const maxBeamWidth, - bool const useMedusa, mpi::MpiComm const& commSession, int const root) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start bcast outputs of decoder from rank %d", root); - - auto request1 = commSession.bcastAsync(*decoderOutputBuffers.newOutputTokensHost, root); - auto request2 = commSession.bcastAsync(*decoderOutputBuffers.finishedSumHost, root); - auto request3 = commSession.bcastAsync(*decoderOutputBuffers.sequenceLengthsHost, root); - auto request4 = returnLogProbs ? commSession.bcastAsync(*decoderOutputBuffers.cumLogProbsHost, root) : nullptr; - auto request5 = returnLogProbs ? commSession.bcastAsync(*decoderOutputBuffers.logProbsHost, root) : nullptr; - auto request6 - = maxBeamWidth > 1 ? commSession.bcastAsync(*decoderState.getCacheIndirectionOutput(), root) : nullptr; - auto request7 = useMedusa ? commSession.bcastAsync(*decoderState.getAcceptedLengthsCumSum(), root) : nullptr; - auto request8 = useMedusa ? commSession.bcastAsync(*decoderState.getAcceptedPackedPaths(), root) : nullptr; - auto request9 = commSession.bcastAsync(*decoderOutputBuffers.finishReasonsHost, root); - - request1->wait(); - request2->wait(); - request3->wait(); - if (request4) - request4->wait(); - if (request5) - request5->wait(); - if (request6) - request6->wait(); - if (request7) - request7->wait(); - if (request8) - request8->wait(); - request9->wait(); - - TLLM_LOG_DEBUG("end bcast outputs of decoder from rank %d", root); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -DecoderSlotAsyncSend::DecoderSlotAsyncSend(TensorPtr const& outputIds, TensorPtr const& sequenceLengths, - TensorPtr const& cumLogProbs, TensorPtr const& logProbs, bool const returnLogProbs, mpi::MpiComm const& commSession, - int const peer) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start send outputs of SlotDecoderBuffers to rank %d", peer); - - mRequest1 = commSession.sendAsync(*outputIds, peer, mpi::MpiTag::kDecoderSlotOutputIds); - mRequest2 = commSession.sendAsync(*sequenceLengths, peer, mpi::MpiTag::kDecoderSlotSequenceLengths); - mRequest3 - = returnLogProbs ? commSession.sendAsync(*cumLogProbs, peer, mpi::MpiTag::kDecoderSlotCumLogProbs) : nullptr; - mRequest4 = returnLogProbs ? commSession.sendAsync(*logProbs, peer, mpi::MpiTag::kDecoderSlotLogProbs) : nullptr; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -DecoderSlotAsyncSend::DecoderSlotAsyncSend(SlotDecoderBuffers const& slotDecoderBuffers, bool const returnLogProbs, - mpi::MpiComm const& commSession, int const peer) - : DecoderSlotAsyncSend(slotDecoderBuffers.outputIds, slotDecoderBuffers.sequenceLengths, - slotDecoderBuffers.cumLogProbs, slotDecoderBuffers.logProbs, returnLogProbs, commSession, peer) -{ -} - -void DecoderSlotAsyncSend::recv(SlotDecoderBuffers const& slotDecoderBuffers, bool const returnLogProbs, - mpi::MpiComm const& commSession, int const peer) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start recv outputs of SlotDecoderBuffers from rank %d", peer); - - commSession.recv(*slotDecoderBuffers.outputIds, peer, mpi::MpiTag::kDecoderSlotOutputIds); - commSession.recv(*slotDecoderBuffers.sequenceLengths, peer, mpi::MpiTag::kDecoderSlotSequenceLengths); - if (returnLogProbs) - { - commSession.recv(*slotDecoderBuffers.cumLogProbs, peer, mpi::MpiTag::kDecoderSlotCumLogProbs); - commSession.recv(*slotDecoderBuffers.logProbs, peer, mpi::MpiTag::kDecoderSlotLogProbs); - } - - TLLM_LOG_DEBUG("end recv outputs of SlotDecoderBuffers from rank %d", peer); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -DecoderSlotAsyncSend::~DecoderSlotAsyncSend() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mRequest1->wait(); - mRequest2->wait(); - if (mRequest3) - mRequest3->wait(); - if (mRequest4) - mRequest4->wait(); - - TLLM_LOG_DEBUG("end send outputs of SlotDecoderBuffers"); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -SlotDecoderBuffers::SlotDecoderBuffers(SizeType32 maxBeamWidth, SizeType32 maxSeqLen, BufferManager const& manager) -{ - outputIds = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kINT32); - outputIdsHost - = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kINT32); - - sequenceLengths = manager.gpu(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kINT32); - sequenceLengthsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kINT32); - - cumLogProbs = manager.gpu(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); - cumLogProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); - - logProbs = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); - logProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp index fa0f69c79235..4d226c4ed3e7 100644 --- a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp +++ b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp @@ -17,7 +17,6 @@ #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/executor/serializeUtils.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" namespace tensorrt_llm::batch_manager { diff --git a/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp b/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp deleted file mode 100644 index 32935e683b83..000000000000 --- a/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp +++ /dev/null @@ -1,56 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/batch_manager/medusaBuffers.h" -#include "tensorrt_llm/runtime/bufferManager.h" - -namespace tensorrt_llm::batch_manager -{ - -void MedusaBuffers::reshape(SizeType32 /* numCtxSequences */, SizeType32 numGenSequences, SizeType32 tokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto attentionPackedMaskShape = attentionPackedMaskDevice->getShape(); - attentionPackedMaskShape.d[0] = numGenSequences * tokensPerStep; - attentionPackedMaskDevice->reshape(attentionPackedMaskShape); - - auto medusaGenerationLengthsShape = medusaGenerationLengthsDevice->getShape(); - medusaGenerationLengthsShape.d[0] = numGenSequences; - medusaGenerationLengthsDevice->reshape(medusaGenerationLengthsShape); - - auto medusaPositionOffsetsShape = medusaPositionOffsetsDevice->getShape(); - medusaPositionOffsetsShape.d[0] = numGenSequences; - medusaPositionOffsetsDevice->reshape(medusaPositionOffsetsShape); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void MedusaBuffers::insertInputTensors( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::WorldConfig const& worldConfig) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - inputBuffers.insert_or_assign("spec_decoding_packed_mask", attentionPackedMaskDevice); - inputBuffers.insert_or_assign("spec_decoding_generation_lengths", medusaGenerationLengthsDevice); - inputBuffers.insert_or_assign("spec_decoding_position_offsets", medusaPositionOffsetsDevice); - inputBuffers.insert_or_assign("spec_decoding_use", medusaUseSpecDecoding); - if (worldConfig.isLastPipelineParallelRank()) - { - outputBuffers.insert_or_assign("medusa_logits", medusaLogitsDevice); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/executor/samplingConfig.cpp b/cpp/tensorrt_llm/executor/samplingConfig.cpp index 3e0ef63f9297..14b2a8b5aef1 100644 --- a/cpp/tensorrt_llm/executor/samplingConfig.cpp +++ b/cpp/tensorrt_llm/executor/samplingConfig.cpp @@ -19,7 +19,6 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" namespace tensorrt_llm::executor { @@ -286,7 +285,7 @@ void SamplingConfig::setBeamWidthArray(OptVec const& beamWidthArray) // Checkers SizeType32 SamplingConfig::checkBeamWidth(SizeType32 beamWidth) { - TLLM_CHECK(beamWidth > 0 && beamWidth <= static_cast(tensorrt_llm::kernels::kMaxBeamWidth)); + TLLM_CHECK(beamWidth > 0 && beamWidth <= kMaxBeamWidth); return beamWidth; } @@ -436,10 +435,10 @@ std::pair const&, SizeType32 const> const SamplingConfig::che if (beamWidthArray.has_value()) { auto array = beamWidthArray.value(); - TLLM_CHECK(array.size() <= static_cast(tensorrt_llm::kernels::kMaxBeamWidthArrayLength)); + TLLM_CHECK(static_cast(array.size()) <= kMaxBeamWidthArrayLength); for (auto const& bm : array) { - TLLM_CHECK(bm > 0 && bm < static_cast(tensorrt_llm::kernels::kMaxBeamWidth)); + TLLM_CHECK(bm > 0 && bm < kMaxBeamWidth); maxBeamWidth = std::max(maxBeamWidth, bm); } } diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels.cu deleted file mode 100644 index 2dadc9d5b682..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels.cu +++ /dev/null @@ -1,372 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" -#include "tensorrt_llm/kernels/beamSearchKernels.h" - -using namespace tensorrt_llm::common; - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -template -void beamSearchKernelLauncher( - T const* logProbs, T const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -#define CASE_K(PBM) \ - { \ - beamSearchKernelLauncher(logProbs, bias, workspace, bh, stream); \ - break; \ - } - -template -void invokeTopkBeamSearch(T const* logProbs, T const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream) -{ - int const nPadBeamWidth{padToNextPowerOfTwo(bh.nBeamWidth)}; - - // case X means X/2 < beam_width <= X - if constexpr (IS_V2) - { - switch (nPadBeamWidth) - { - case 1: - case 2: - case 4: CASE_K(4) - case 8: CASE_K(8) - case 16: CASE_K(16) -#ifndef FAST_BUILD // Skip beam width > 16 - case 32: CASE_K(32) - case 64: CASE_K(64) - case 128: CASE_K(128) - case 256: CASE_K(256) - case 512: CASE_K(512) - case 1024: CASE_K(1024) -#endif // FAST_BUILD - } - } - else // V1, only use kernels of `beam_width <= kMaxBeamWidthForV1` - { - switch (nPadBeamWidth) - { - case 1: - case 2: - case 4: CASE_K(4) - case 8: CASE_K(8) - } - } -} - -#undef CASE_K - -template void invokeTopkBeamSearch( - float const* logProbs, float const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -template void invokeTopkBeamSearch( - float const* logProbs, float const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -template void invokeTopkBeamSearch( - half const* logProbs, half const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -template void invokeTopkBeamSearch( - half const* logProbs, half const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -__global__ void updateCacheIndirectionKernel( - int* tgtCI, int const* srcCI, BeamHypotheses bh, int const nMaxAttentionWindow, int const nSinkTokenLength) -{ - // Update cache indirections which steps are between `bh.inputLength[x]` to `sequenceLengths[x]` - int const step = blockIdx.x * blockDim.x + threadIdx.x; - size_t const nBM{bh.nBeamWidth}; - size_t const nBMIn{bh.nBeamWidthIn}; - size_t const nBMOut{bh.nBeamWidthOut}; - size_t const nMSL{bh.nMaxSeqLen}; - int const indexBatch = blockIdx.y; - int const batchSlot = bh.batchSlots[indexBatch]; - int const tgtIndexBeam = blockIdx.z; - int const tgtIndexBatchBeam = batchSlot * nBM + tgtIndexBeam; - int const lastStep{bh.sequenceLengths[tgtIndexBatchBeam] - 1}; // minus 1 since it is updated in stage 3 kernel - - // Return early when at least one of the conditions is true: - // 1. `step` is out of the bound - // 2. `step` is inside of input part (since context KV Cache is shared) - // 3. `step` is outside of attention widow - if (step >= nMSL || step < bh.inputLengths[tgtIndexBatchBeam] || step < (nMSL - nMaxAttentionWindow)) - { - return; - } - - // Keep all past tokens by parentIdsPtr - int const srcIndexBeam = bh.parentIdsPtr[batchSlot][tgtIndexBeam * nMSL + lastStep]; - // Return early when the source beam isfinished - if (bh.finished[tgtIndexBatchBeam].isFinished()) - { - return; - } - - int const stepCirc = (step >= nSinkTokenLength) - ? nSinkTokenLength + (step - nSinkTokenLength) % (nMaxAttentionWindow - nSinkTokenLength) - : step; - // Consider cyclic kv cache for the indir tables - uint32_t const tgtOffset = batchSlot * nBMOut * nMaxAttentionWindow + tgtIndexBeam * nMaxAttentionWindow + stepCirc; - uint32_t const srcOffset = batchSlot * nBMIn * nMaxAttentionWindow + srcIndexBeam * nMaxAttentionWindow + stepCirc; - tgtCI[tgtOffset] = (step == lastStep) ? tgtIndexBeam : srcCI[srcOffset]; -} - -void invokeUpdateCacheIndirection(int* tgtCI, int const* srcCI, BeamHypotheses& bh, - runtime::SizeType32 const maxAttentionWindow, runtime::SizeType32 sinkTokenLength, cudaStream_t stream) -{ - dim3 const grid(common::roundUp(bh.nMaxSeqLen, 32), bh.nBatchSize, bh.nBeamWidthOut); - updateCacheIndirectionKernel<<>>(tgtCI, srcCI, bh, maxAttentionWindow, sinkTokenLength); - sync_check_cuda_error(stream); -} - -__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, FinishedState const* finished, int const* endIds, float const* diversityRates, - runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM) -{ - int const bid = blockIdx.x; // Index of request in batch - runtime::SizeType32 const slot = batchSlots[bid]; - float const diversityRate{diversityRates[slot]}; - float* pLocalLogProbs = pStage1LogProbs + bid * nBMIn * nBMOut * 2; - int const* pLocalIds = pStage1Ids + bid * nBMIn * nBMOut * 2; - - for (int i = threadIdx.x; i < nBMIn * nBMOut * 2; i += blockDim.x) - { - int const iBMIn = i / (nBMOut * 2); - if (finished[slot * nBM + iBMIn].isFinished()) - { - // In V2 path, i is a candidate-slot index (0..nBMIn*nBMOut*2-1), NOT a vocab token id. - // Use pStage1Ids to look up the actual token id for the EOS comparison. - bool const isEOS = (pLocalIds[i] == endIds[slot]); - // Keep only the EOS candidate with its proper cumulative score; suppress all others. - pLocalLogProbs[i] = isEOS ? (pLocalLogProbs[i] + cumLogProbs[slot * nBM + iBMIn]) : -FLT_MAX; - } - else - { - // nBM is used in VBWS since `cumLogProbs` is initialized with kMaxBeamWidth earlier than BeamSearchLayer - pLocalLogProbs[i] += cumLogProbs[slot * nBM + iBMIn] + diversityRate * iBMIn; - } - } - return; -} - -__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, FinishedState const* finished, int const* endIds, float const* diversityRates, - runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM) -{ - int const bid = blockIdx.x; // Index of request in batch - runtime::SizeType32 const slot = batchSlots[bid]; - float const diversityRate{diversityRates[slot]}; - half* pLocalLogProbs = pStage1LogProbs + bid * nBMIn * nBMOut * 2; - int const* pLocalIds = pStage1Ids + bid * nBMIn * nBMOut * 2; - - for (int i = threadIdx.x; i < nBMIn * nBMOut * 2; i += blockDim.x) - { - int const iBMIn = i / (nBMOut * 2); - if (finished[slot * nBM + iBMIn].isFinished()) - { - // In V2 path, i is a candidate-slot index (0..nBMIn*nBMOut*2-1), NOT a vocab token id. - // Use pStage1Ids to look up the actual token id for the EOS comparison. - bool const isEOS = (pLocalIds[i] == endIds[slot]); - // Keep only the EOS candidate with its proper cumulative score; suppress all others. - pLocalLogProbs[i] - = isEOS ? (half) (float(pLocalLogProbs[i]) + cumLogProbs[slot * nBM + iBMIn]) : (half) -HALF_FLT_MAX; - } - else - { - // nBM is used in VBWS since `cumLogProbs` is initialized with kMaxBeamWidth earlier than BeamSearchLayer - pLocalLogProbs[i] += cumLogProbs[slot * nBM + iBMIn] + diversityRate * iBMIn; - } - } - return; -} - -__global__ void gatherId(int const* __restrict pStage1Id, int* __restrict pStage2Id, size_t const nBS, - size_t const nBMIn, size_t const nBMOut, size_t const nV) -{ - // Use topK output `pStage1Id` and `pStage1Id` to get the index of a new token in `logProbs` for each beam. - // - // clang-format off - // - // Example for normal beam search: - // nBS = 3, nBM = 2, nV = 5, use logProbs with integer values here for simplicity. - // ┏┏ 46 35 47 18 67 ┓┓ ┏┏ 67 47 46 35 ┓┓ ┏┏ 4 2 0 1 ┓┓ - // ┃┗ 76 23 74 73 17 ┛┃ ┃┗ 76 74 73 23 ┛┃ ┃┗ 0 2 3 1 ┛┃ - // ┃┏ 67 49 98 88 74 ┓┃ A ┃┏ 98 88 74 67 ┓┃ ┃┏ 2 3 4 0 ┓┃ C ┏ 76 74 73 67 ┓ ┏ 4 5 6 0 ┓ D ┏ 5 7 8 4 ┓ - // ┃┗ 12 70 77 22 88 ┛┃ ---> ┃┗ 88 77 70 22 ┛┃ ┃┗ 4 2 1 3 ┛┃ ---> ┃ 98 88 88 77 ┃ ┃ 0 1 4 5 ┃ ---> ┃ 2 3 9 7 ┃ - // ┃┏ 55 15 72 3 84 ┓┃ ┃┏ 74 72 55 15 ┓┃ ┃┏ 4 2 0 1 ┓┃ ┗ 98 93 84 77 ┛ ┗ 4 5 0 6 ┛ ┗ 9 6 4 5 ┛ - // ┗┗ 77 93 14 60 98 ┛┛ ┗┗ 98 93 77 60 ┛┛ ┗┗ 4 1 0 3 ┛┛ - // logProbs stage1LogProbs stage1Id stage2LogProbs stage2Id output-stage2Id - // - // For `stage2LogProbs[2][3] == 77`, - // original batch index in logProbs: blockIdx.x -> 2 (a) - // original beam index in logProbs: stage2Id[2][3] / (nBM * 2) -> 1 (b) - // row index in stage1Probs: a * nBM + b -> 5 (c) - // column index in stage1*: stage2Id[2][3] % (nBM * 2) -> 2 (d) - // column index in logProbs: stage1Id[c][d] -> 0 (e) - // pad for previous tokens: b * nV -> 5 (f) - // final output: e + f -> 5 - // - // ======================================================================================================== - // Example for VBWS: - // nBS = 2, nBMIn = 3, nBMOut = 5, nBM = 7, nV = 11, use logProbs with integer values here for simplicity. - // ┏┏ 46 35 47 18 67 76 23 74 73 17 67 ┓┓ ┏┏ 76 74 73 67 67 47 46 35 23 18 ┓┓ ┏┏ 5 7 8 4 10 2 0 1 6 3 ┓┓ - // ┃┃ 49 98 88 74 12 70 77 22 88 55 15 ┃┃ ┃┃ 98 88 88 77 74 70 55 49 22 15 ┃┃ ┃┃ 1 2 8 6 3 5 9 0 7 10 ┃┃ - // ┃┗ 72 3 84 77 93 14 60 98 65 4 20 ┛┃ A ┃┗ 98 93 84 77 72 65 60 20 14 4 ┛┃ ┃┗ 7 4 2 3 0 8 6 10 5 9 ┛┃ C - // ┃┏ 16 34 71 38 19 91 5 81 97 43 79 ┓┃ ---> ┃┏ 97 91 81 79 71 43 38 34 19 16 ┓┃ ┃┏ 8 5 7 10 2 9 3 1 4 0 ┓┃ ---> - // ┃┃ 2 22 77 37 57 33 57 41 27 73 88 ┃┃ ┃┃ 88 77 73 57 57 41 37 33 27 22 ┃┃ ┃┃ 10 2 9 4 6 7 3 5 8 1 ┃┃ - // ┗┗ 77 16 23 22 82 89 6 77 67 15 31 ┛┛ ┗┗ 89 82 77 77 67 31 23 22 16 15 ┛┛ ┗┗ 5 4 0 7 8 10 2 3 1 9 ┛┛ - // logProbs stage1LogProbs stage1Id - // - // C ┏ 98 98 93 88 88 84 77 77 76 74 ┓ ┏ 10 20 21 11 12 22 13 23 0 1 ┓ D ┏ 12 29 26 13 19 24 17 25 5 7 ┓ - // ---> ┗ 97 91 89 88 82 81 79 77 77 77 ┛ ┗ 0 1 20 10 21 2 3 11 22 23 ┛ ---> ┗ 8 5 27 21 26 7 10 13 22 29 ┛ - // stage2LogProbs stage2Id output-stage2Id - // - // For `stage2LogProbs[1][4] == 82`, - // original batch index in logProbs: blockIdx.x -> 1 (a) - // original beam index in logProbs: stage2Id[1][4] / (nBMOut * 2) -> 2 (b) - // row index in stage1LogProbs: a * nBMIn + b -> 5 (c) - // column index in stage1*: stage2Id[1][4] % (nBMOut * 2) -> 1 (d) - // column index in logProbs: stage1Id[c][d] -> 4 (e) - // pad for previous tokens: b * nV -> 22 (f) - // final output: e + f -> 26 output-stage2Id[1][4] - // - // clang-format on - int const a = blockIdx.x; // Index of request in batch - for (int j = threadIdx.x; j < nBMOut * 2; j += blockDim.x) - { - int const index = a * (nBMOut * 2) + j; - int const stage2Id = pStage2Id[index]; - int const b = stage2Id / (nBMOut * 2); - int const c = a * nBMIn + b; - int const d = stage2Id % (nBMOut * 2); - int const e = pStage1Id[c * (nBMOut * 2) + d]; - int const f = b * nV; - pStage2Id[index] = e + f; - } - return; -} - -void BeamHypotheses::print() -{ -#if BEAM_SEARCH_DEBUG - cudaDeviceSynchronize(); - printf("================ print BeamHypotheses start\n"); - - PRINT(this->bReturnNormedScore); - PRINT(this->bVBWS); - PRINT(this->nMaxBatchSize); - PRINT(this->nBatchSize); - PRINT(this->nBeamWidth); - PRINT(this->nBeamWidthIn); - PRINT(this->nBeamWidthOut); - PRINT(this->nMaxSeqLen); - PRINT(this->nVocabSize); - PRINT(this->nVPart); - PRINT(this->nByteMaxSharedMemoryPerBlock); - PRINT(this->nByteSharedMemoryStage1); - PRINT(this->nByteSharedMemoryStage3); - size_t const mbs = this->nMaxBatchSize; - size_t const nbs = this->nBatchSize; - size_t const nbm = this->nBeamWidth; - size_t const nbmo = this->nBeamWidthOut; - size_t const msl = this->nMaxSeqLen; - - PH2(this->diversityRates, nbs); - PH2(this->lengthPenalties, nbs); - PH2(this->earlyStoppings, nbs); - PH3(this->beamWidthArraysHost, nbs * kMaxBeamWidthArrayLength, kMaxBeamWidthArrayLength); - PH2(this->nBeamWidthInHost, nbs); - PH2(this->nBeamWidthOutHost, nbs); - - PH2(this->inputLengths, nbs * nbm); - PH2(this->endIds, nbs); - PH2(this->batchSlots, nbs); - - PH3(this->outputIds, nbs * nbm * msl, msl); - PH3(this->logProbs, nbs * nbm * msl, msl); - PH3(this->sequenceLengths, nbs * nbm, nbm); - PH3(this->cumLogProbs, nbs * nbm, nbm); - - PH3(this->outputIdsCBA, mbs * nbmo * 2 * msl, msl); - PH3(this->logProbsCBA, mbs * nbmo * 2 * msl, msl); - PH3(this->sequenceLengthsCBA, mbs * nbmo * 2, nbmo * 2); - PH3(this->cumLogProbsCBA, mbs * nbmo * 2, nbmo * 2); - PH3(this->normedScoresCBA, mbs * nbmo * 2, nbmo * 2); - PH2(this->numBeamsCBA, mbs); - PH2(this->minNormedScoresCBA, mbs); - - // PH2(this->batchDones, nbs); - uint8_t* finished = reinterpret_cast(this->finished); - PH2(finished, nbs * nbm); - - std::vector batchSlots(nbs, 0); - cudaMemcpy(batchSlots.data(), this->batchSlots, sizeof(runtime::SizeType32) * nbs, cudaMemcpyDeviceToHost); - - std::vector outputIdsPtr(nbs, 0); - cudaMemcpy(outputIdsPtr.data(), this->outputIdsPtr, sizeof(int*) * nbs, cudaMemcpyDeviceToHost); - - std::vector parentIdsPtr(nbs, 0); - cudaMemcpy(parentIdsPtr.data(), this->parentIdsPtr, sizeof(int*) * nbs, cudaMemcpyDeviceToHost); - cudaDeviceSynchronize(); - - for (int i = 0; i < nbs; ++i) - { - int slot = batchSlots[i]; - printf("slot=%d\n", slot); - printf("outputIdsPtr[slot]=%p\n", outputIdsPtr[slot]); - PH3(outputIdsPtr[slot], nbm * msl, msl); - } - for (int i = 0; i < nbs; ++i) - { - int slot = batchSlots[i]; - printf("slot=%d\n", slot); - printf("parentIdsPtr[slot]=%p\n", parentIdsPtr[slot]); - PH3(parentIdsPtr[slot], nbm * msl, msl); - } - - // May not available in some context - // PH3(this->outputIdsUnfinish, nbs * nbm * msl, msl); - // PH3(this->parentIdsUnfinish, nbs * nbm * msl, msl); - - printf("================ print BeamHypotheses stop\n"); -#endif -} - -template -void printLogProbs(T const* x, int const nBS, int const nBMIn, int const nBM, int const nV) -{ - for (int bs = 0; bs < nBS; ++bs) - { - T const* ptrBatch = x + bs * nBM * nV; - printArrayInfo(ptrBatch, nBMIn * nV, std::string("Request ") + std::to_string(bs)); - for (int bm = 0; bm < nBMIn; ++bm) - { - T const* ptrBeam = ptrBatch + bm * nV; - printArrayInfo(ptrBeam, nV, std::string("Beam ") + std::to_string(bm), true); - } - } -} - -template void printLogProbs(float const* x, int const nBS, int const nBMIn, int const nBM, int const nV); -template void printLogProbs(half const* x, int const nBS, int const nBMIn, int const nBM, int const nV); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels.h b/cpp/tensorrt_llm/kernels/beamSearchKernels.h deleted file mode 100644 index 345e4659c941..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels.h +++ /dev/null @@ -1,226 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/topkLastDim.h" // Air TopK -#include "tensorrt_llm/runtime/common.h" - -#define BEAM_SEARCH_DEBUG 0 - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -static size_t constexpr kMaxBeamWidth = 1024; // Max beam width supported in TRT-LLM now -static size_t constexpr kMaxBeamWidthForV1 = 8; // Max beam width for V1 workflow (V2 for larger) -static size_t constexpr kMaxBeamWidthArrayLength = 8; // Max length of beam width array of a request -static size_t constexpr kThreadForSmallBeamWidth = 256; // Max count of thread for stage 1 in V1 workflow -static size_t constexpr kMaxVPartStage1 = 128; // Max vocab part count for stage 1 in V1 workflow - -struct BeamHypotheses -{ - // clang-format off - // MBS: max_batch_size, BS: batch_size, BM: beam_width, MSL: max_seq_length - // %%: parameter name in file generation.py (python workflow) - // Candidate beams: a beam which generates end_id or its sequence length reaches MSL - // Candidate-Beam-Array (CBA): The arrays to place the candidate beams and related information - // Variable-Beam-Width-Search (VBWS): A search mode that allows using different beam width for each step - - // Scalar values - bool bReturnNormedScore{false}; // Return `normedScore` or `cumLogProbs`, always be `false` now - bool bVBWS{false}; // whether to use VBWS for Beam-Search - size_t nMaxBatchSize{0}; // Buildtime max batch size - size_t nBatchSize{0}; // Runtime batch size - size_t nBeamWidth{0}; // Runtime beam width - size_t nBeamWidthIn{0}; // Scalar value of current input beam width, for VBWS - size_t nBeamWidthOut{0}; // Scalar value of current output beam width, for VBWS - size_t nMaxSeqLen{0}; // - size_t nVocabSize{0}; // Vocab Size Padded - size_t nVPart{0}; // Count of vocab_size_padded divided - size_t nByteMaxSharedMemoryPerBlock{0}; // Device information - size_t nByteSharedMemoryStage1{0}; // Dynamic shared memory size of stage 1 - size_t nByteSharedMemoryStage3{0}; // Static shared memory size of stage 3 - - // Pointers from SamplingConfig - float const* diversityRates{nullptr}; // [BS] - float const* lengthPenalties{nullptr}; // [BS] - int const* earlyStoppings{nullptr}; // [BS] - int const* beamWidthArraysHost{nullptr}; // [BS, kMaxBeamWidthArrayLength] for VBWS - int* nBeamWidthInHost{nullptr}; // [BS], cpu for VBWS, beam width of last forward computation - int* nBeamWidthOutHost{nullptr}; // [BS], cpu for VBWS, beam width of next forward computation - - // Pointers from input - int const* inputLengths{nullptr}; // [BS, BM] %% context_length - int const* endIds{nullptr}; // [BS, BM] %% self.end_ids - runtime::SizeType32 const* batchSlots{nullptr}; // [BS] - - // Pointers for output - int* outputIds{nullptr}; // [BS, BM, MSL] %% self.output_ids only used in gather_tree - float* logProbs{nullptr}; // [BS, BM, MSL] %% self.log_probs only used in gather_tree - float* logProbsTiled{nullptr}; // [MSL, MBS, BM] %% self.log_probs_tiled - int* sequenceLengths{nullptr}; // [BS, BM] %% self.sequence_length_buffer - float* cumLogProbs{nullptr}; // [BS, BM] %% self.cum_log_probs - - // Pointers of CBA - int* outputIdsCBA{nullptr}; // [BS, BM*2, MSL] %% self.beam_hyps_output_ids_cba - float* logProbsCBA{nullptr}; // [BS, BM*2, MSL] %% self.beam_hyps_log_probs_cba - int* sequenceLengthsCBA{nullptr}; // [BS, BM*2] %% self.beam_hyps_seq_len_cba - float* cumLogProbsCBA{nullptr}; // [BS, BM*2] %% self.beam_hyps_cum_log_probs_cba - float* normedScoresCBA{nullptr}; // [BS, BM*2] %% self.beam_hyps_normed_scores_cba - int* numBeamsCBA{nullptr}; // [BS] %% self.beam_hyps_num_beams number of beams in CBA - float* minNormedScoresCBA{nullptr}; // [BS] %% self.beam_hyps_min_normed_scores worst score in CBA - - // Pointers related to beam search process, they are initialized in those two functions: - // [gptDecoder.cpp] GptDecoder::forward or [dynamicDecodeOp.cpp] FtDynamicDecode::forward - bool* batchDones{nullptr}; // [BS] %% self.beam_hyps_is_done whether a whole batch is finished - ::tensorrt_llm::kernels::FinishedState* finished{nullptr}; // [BS*BM], uint8 %% self.finished whether and how a beam is finished - - // Pointers for backtrack of the beams, they are relocated in [dynamicDecodeLayer.cpp] DynamicDecodeLayer::prepareIdsPtrs - int** outputIdsPtr{nullptr}; // [BS][BM, MSL] %% self.output_ids - int** parentIdsPtr{nullptr}; // [BS][BM, MSL] %% self.parent_ids - - // Pointers for gather_tree(), read the unfinished beams from them and write to CBA for the final selection - int const* outputIdsUnfinish{nullptr}; // [BS, BM, MSL] %% self.output_ids - int const* parentIdsUnfinish{nullptr}; // [BS, BM, MSL] %% self.parent_ids - - // clang-format on - - void print(); -}; - -__inline__ int padToNextPowerOfTwo(int const n) -{ - // Pad n up to the nearest power of 2 - int recursor = n - 1; - int res = 2; - while (recursor >>= 1) - res <<= 1; - return res; -} - -template -__device__ __forceinline__ T applyLengthPenalty(T const log_prob, int const length, float const length_penalty) -{ - // score = log(prob) / (length ^ length_penalty) - if (length_penalty == 0.0f || length == 1) - { - return log_prob; - } - return log_prob / static_cast(powf(static_cast(length), length_penalty)); -} - -template -void invokeTopkBeamSearch(T const* logProbs, T const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -void invokeUpdateCacheIndirection(int* tgtCI, int const* srcCI, BeamHypotheses& bh, - runtime::SizeType32 const maxAttentionWindow, runtime::SizeType32 sinkTokenLength, cudaStream_t stream); - -__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, - float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, - size_t const nBMOut, size_t const nBM); - -__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, - float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, - size_t const nBMOut, size_t const nBM); - -__global__ void gatherId(int const* __restrict pStage1Id, int* __restrict pStage2Id, size_t const nBS, - size_t const nBMIn, size_t const nBMOut, size_t const nV); - -void printLogProbs(float const* x, int const nBS, int const nBMIn, int const nBM, int const nV); - -// for Beam Search debug -#if BEAM_SEARCH_DEBUG -#define BID 0 - -#define LINE(x) printf(x "@L%d\n", __LINE__); - -#define PRINT(x) \ - { \ - printf(#x "="); \ - print_element_(x); \ - printf("\n"); \ - } - -// Host function -#define PRINT_HOST(x, nRow, nCol, nColPadded) \ - { \ - if (x == nullptr) \ - { \ - printf(#x "=nullptr\n"); \ - } \ - else \ - { \ - printf(#x "=\n"); \ - printMatrix(x, nRow, nCol, nColPadded); \ - } \ - } -#define PH2(x, nCol) PRINT_HOST(x, 1, nCol, nCol) -#define PH3(x, nElement, nCol) PRINT_HOST(x, ((nElement) / (nCol)), nCol, nCol) - -// Device function -#define PRINT_DEVICE(x, nRow, nCol, nColPadded) \ - { \ - if (x == nullptr) \ - { \ - printf(#x "=nullptr\n"); \ - } \ - else \ - { \ - printf(#x "=\n"); \ - printMatrixDevice(x, nRow, nCol, nColPadded); \ - } \ - } -#define PD2(x, nCol) PRINT_DEVICE(x, 1, nCol, nCol) -#define PD3(x, nElement, nCol) PRINT_DEVICE(x, ((nElement) / (nCol)), nCol, nCol) - -// Device function -#define WITH(blockIdxx, bSync, code) \ - { \ - if (bSync) \ - { \ - __syncthreads(); \ - } \ - if (blockIdx.x == (blockIdxx) && blockIdx.y == 0 && blockIdx.z == 0 && threadIdx.x == 0 && threadIdx.y == 0 \ - && threadIdx.z == 0) \ - { \ - code \ - } \ - if (bSync) \ - { \ - __syncthreads(); \ - } \ - } - -#else -#define LINE(x) -#define PRINT(x) -#define QH(x, y, z, w) -#define PH2(x, nCol) -#define PH3(x, nElement, nCol) -#define PRINT_DEVICE(x, y, z, w) -#define PD2(x, nCol) -#define PD3(x, nElement, nCol) -#define WITH(x, y, z) -#endif - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels1024.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels1024.cu deleted file mode 100644 index 4d6005558588..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels1024.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#ifndef FAST_BUILD // Skip beam_width larger than 16 -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 1024, true); -INSTANTIATE_BEAM_SEARCH(half, 1024, true); -#endif // FAST_BUILD - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels128.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels128.cu deleted file mode 100644 index bf23a844b927..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels128.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#ifndef FAST_BUILD // Skip beam_width larger than 16 -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 128, true); -INSTANTIATE_BEAM_SEARCH(half, 128, true); -#endif // FAST_BUILD - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels16.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels16.cu deleted file mode 100644 index 50bf27b14273..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels16.cu +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 16, true); -INSTANTIATE_BEAM_SEARCH(half, 16, true); -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels256.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels256.cu deleted file mode 100644 index fae7cd927e91..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels256.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#ifndef FAST_BUILD // Skip beam_width larger than 16 -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 256, true); -INSTANTIATE_BEAM_SEARCH(half, 256, true); -#endif // FAST_BUILD - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels32.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels32.cu deleted file mode 100644 index d414d268c0db..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels32.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#ifndef FAST_BUILD // Skip beam_width larger than 16 -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 32, true); -INSTANTIATE_BEAM_SEARCH(half, 32, true); -#endif // FAST_BUILD - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels4.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels4.cu deleted file mode 100644 index d1815d85e3e0..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels4.cu +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -INSTANTIATE_BEAM_SEARCH(float, 4, false); -INSTANTIATE_BEAM_SEARCH(float, 4, true); -INSTANTIATE_BEAM_SEARCH(half, 4, false); -INSTANTIATE_BEAM_SEARCH(half, 4, true); -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels512.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels512.cu deleted file mode 100644 index 005f44e5e755..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels512.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#ifndef FAST_BUILD // Skip beam_width larger than 16 -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 512, true); -INSTANTIATE_BEAM_SEARCH(half, 512, true); -#endif // FAST_BUILD - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels64.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels64.cu deleted file mode 100644 index 87a34b2d07e7..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels64.cu +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#ifndef FAST_BUILD // Skip beam_width larger than 16 -// Skip V1 kernels if beam_width > kMaxBeamWidthForV1 -INSTANTIATE_BEAM_SEARCH(float, 64, true); -INSTANTIATE_BEAM_SEARCH(half, 64, true); -#endif // FAST_BUILD - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels8.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels8.cu deleted file mode 100644 index 7b84b37050ae..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernels8.cu +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchKernelsTemplate.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -INSTANTIATE_BEAM_SEARCH(float, 8, false); -INSTANTIATE_BEAM_SEARCH(float, 8, true); -INSTANTIATE_BEAM_SEARCH(half, 8, false); -INSTANTIATE_BEAM_SEARCH(half, 8, true); -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h deleted file mode 100644 index 09b114ef9752..000000000000 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h +++ /dev/null @@ -1,768 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#ifndef CUDART_VERSION -#error CUDART_VERSION Undefined! -#elif (CUDART_VERSION >= 11050) -#include - -#else -#include "3rdparty/cub/cub.cuh" -#endif - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" -#include "tensorrt_llm/common/stringUtils.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/kernels/decodingCommon.h" - -using namespace tensorrt_llm::common; - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -#pragma nv_diag_suppress static_var_with_dynamic_init - -template -__launch_bounds__(BLOCK_SIZE) __global__ void beamStage1Kernel(T const* __restrict logProbs, T const* __restrict bias, - float* __restrict pStage3, int const* __restrict endIds, FinishedState const* __restrict finished, int const nV, - runtime::SizeType32 const* batchSlots) -{ - int const nBM = gridDim.y; - int const tid = threadIdx.x; - int const slot = batchSlots[blockIdx.x]; - int const nVLocal = (nV + gridDim.z - 1) / gridDim.z; - int const indexLeft = nVLocal * blockIdx.z; - int const indexRight = std::min(indexLeft + nVLocal, nV); - int const nVOffset = (blockIdx.x * nBM + blockIdx.y) * nV; - int const nVChunk = indexRight - indexLeft; - T const MAX_T_VAL = std::is_same_v ? HALF_FLT_MAX : FLT_MAX; - - using KVPair = cub::KeyValuePair; - using BlockReduceTopK = cub::BlockReduce; - cub::ArgMax argmax; - - __shared__ float smemOutput[PBM * 4]; - __shared__ int threadToUpdate; - __shared__ typename BlockReduceTopK::TempStorage smemReduceBuffer; - extern __shared__ char smem[]; - T* smemLogProbs = reinterpret_cast(smem); - - // Load element from logProbs to smemLogProbs and do argmax meanwhile - // Each thread is responsible for `nVLocal / BLOCK_SIZE` elements - // Dynamic shared memory size: sizeof(T) * (nV + nVPart - 1) / nVPart - KVPair kvLocal{-1, -MAX_T_VAL}; - for (int i = indexLeft + tid; i < indexRight; i += BLOCK_SIZE) - { - T const b{bias == nullptr ? (T) 0.0f : bias[i]}; - int const index = i - indexLeft; - T const value = (finished[slot * nBM + blockIdx.y].isFinished()) ? (i == endIds[slot] ? MAX_T_VAL : -MAX_T_VAL) - : (logProbs[nVOffset + i] + b); - smemLogProbs[index] = value; - kvLocal = argmax(kvLocal, {index, value}); - } - __syncthreads(); - - // Search the top 2K elements among `nVLocal` elements of this ThreadBlock and write into smemOutput - for (int i = 0; i < 2 * nBM; ++i) - { - // Pop the element with largest value to "smemOutput" per iteration - KVPair kv = BlockReduceTopK(smemReduceBuffer).Reduce(kvLocal, argmax); - if (tid == 0) - { - int const index = nVOffset + indexLeft + kv.key; - reinterpret_cast(smemOutput)[i] = index; - smemOutput[PBM * 2 + i] = kv.value; - smemLogProbs[kv.key] = -MAX_T_VAL; // Invalidate the value of the popped element - threadToUpdate = kv.key % BLOCK_SIZE; - } - __syncthreads(); - - if (tid == threadToUpdate && i < 2 * nBM - 1) - { - // The thread popped the element need to update its kvLocal - // No need to do this in the last iteration - kvLocal.key = nV - 1; - kvLocal.value = -MAX_T_VAL; - for (int index = tid; index < nVChunk; index += BLOCK_SIZE) - { - kvLocal = argmax(kvLocal, {index, smemLogProbs[index]}); - } - } - __syncthreads(); - } - // Write the smemOutput into pStage3 - pStage3 += (blockIdx.x * nBM + blockIdx.y) * gridDim.z * PBM * 4 + blockIdx.z * PBM * 4; - for (int i = tid; i < PBM * 4; i += BLOCK_SIZE) - { - pStage3[i] = smemOutput[i]; - } -} - -template -__launch_bounds__(BLOCK_SIZE) __global__ - void beamStage2Kernel(int* __restrict pStage2Ids, T* __restrict pStage2LogProbs, float* __restrict pStage3, - float const* __restrict cumLogProbs, runtime::SizeType32 const* batchSlots, int const nV, int const nVPart) -{ - int const nBM = gridDim.y; - int const gbid = blockIdx.x * gridDim.y + blockIdx.y; - int const tid = threadIdx.x; - int const slot = batchSlots[blockIdx.x]; - T const MAX_T_VAL = std::is_same_v ? HALF_FLT_MAX : FLT_MAX; - - using KVPair = cub::KeyValuePair; - using BlockReduceTopK = cub::BlockReduce; - cub::ArgMax argmax; - - __shared__ KVPair smemOutput[PBM * 2]; - __shared__ typename BlockReduceTopK::TempStorage smemReduceBuffer; - - // Load data from stage 1 - float* pStage2Temp = pStage3 + PBM * 4 * gbid * nVPart; - if constexpr (IS_FAST) - { - // Use shared memory instead of global memory - extern __shared__ char smem[]; - float* smemVal = reinterpret_cast(smem); - for (int idx = tid; idx < PBM * 4 * nVPart; idx += BLOCK_SIZE) - { - smemVal[idx] = pStage2Temp[idx]; - } - pStage2Temp = smemVal; - __syncthreads(); - } - - // Find the top 2K across all nVPart - for (int k = 0; k < 2 * nBM; ++k) - { - KVPair kvLocal{nV - 1, -MAX_T_VAL}; - if (tid < nVPart) - { - for (int i = 0; i < 2 * nBM; ++i) - { - int const index = tid * PBM * 4 + i; - T const topValue = pStage2Temp[index + PBM * 2]; - kvLocal = argmax(kvLocal, {index, topValue}); - } - } - KVPair kv = BlockReduceTopK(smemReduceBuffer).Reduce(kvLocal, argmax); - if (tid == 0) - { - // Replace local offset into global offset and store kv pairs in shared memory - int const offsetLocal = kv.key; - kv.key = reinterpret_cast(pStage2Temp)[offsetLocal]; - smemOutput[k] = kv; - // Invalidate the maximum value within the chunk - reinterpret_cast(pStage2Temp)[offsetLocal] = nV - 1; // id in shared memory - pStage2Temp[offsetLocal + PBM * 2] = -MAX_T_VAL; // value in shared memory - } - __syncthreads(); - } - if (tid == 0) - { - auto const cumLogProb = cumLogProbs[slot * nBM + blockIdx.y]; - for (int i = 0; i < 2 * nBM; ++i) - { - pStage2Ids[gbid * 2 * nBM + i] = smemOutput[i].key; - pStage2LogProbs[gbid * 2 * nBM + i] = (float) smemOutput[i].value + cumLogProb; - } - } -} - -template -__launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( - int const* __restrict pStage2Ids, T const* __restrict pStage2LogProbs, float* __restrict pStage3, BeamHypotheses bh) -{ - T const MAX_T_VAL = std::is_same_v ? HALF_FLT_MAX : FLT_MAX; - int const bid = blockIdx.x; // Index of Batch - int const tid = threadIdx.x; - int const slot = bh.batchSlots[bid]; - size_t const nMBS{bh.nMaxBatchSize}; // Only for bh.logProbsTiled - size_t const nBM{bh.nBeamWidth}; - // size_t const nBMIn{bh.bVBWS ? bh.nBeamWidthIn : bh.nBeamWidth}; - size_t const nBMOut{bh.bVBWS ? bh.nBeamWidthOut : bh.nBeamWidth}; - size_t const nMSL{bh.nMaxSeqLen}; - size_t const nV{bh.nVocabSize}; - float const diversityRate{bh.diversityRates[slot]}; - float const lengthPenalty{bh.lengthPenalties[slot]}; - int const earlyStopping{bh.earlyStoppings[slot]}; - - using KVPair = cub::KeyValuePair; - __shared__ float smemCumLogProbs[PBM]; - __shared__ int smemSeqLen[PBM]; - __shared__ KVPair smemTopKV[(IS_V2) ? 1 : PBM * 2]; // Just a placeholder in V2 workflow - __shared__ int smemNBeamForNextStep; - - if (bh.numBeamsCBA != nullptr) - { - // Beam search is enabled - if (bh.numBeamsCBA[slot] == 0 && tid == 0) - { - // Initialize worst score in the first call - bh.minNormedScoresCBA[slot] = 0.0f; // logProbs is in range (-inf, 0] - } - else if (earlyStopping == 1 && bh.numBeamsCBA[slot] >= nBM || earlyStopping != 1 && bh.batchDones[slot]) - { - // Condition of early return: - // 1. In EarlyStopping mode, and we have got enough beams - // 2. In NonEarlyStopping mode, and this batch has been marked as done - return; - } - } - - // This TopK is needless in V2 workflow - if constexpr (IS_V2) - { - pStage2Ids += bid * nBMOut * 2; - pStage2LogProbs += bid * nBMOut * 2; - } - else - { - int const nCandidate = nBM * nBM * 2; - pStage2Ids += bid * nCandidate; - pStage2LogProbs += bid * nCandidate; - KVPair kvLocal{nCandidate - 1, -MAX_T_VAL}; - cub::ArgMax argmax; - extern __shared__ char smem[]; - T* smemVal = nullptr; - if constexpr (IS_FAST) - { - smemVal = reinterpret_cast(smem); - } - else - { - smemVal = reinterpret_cast(pStage3); - } - - for (int i = tid; i < nCandidate; i += BLOCK_SIZE) - { - int const index = bh.numBeamsCBA == nullptr ? i % nBM : i / 2 / nBM; - T const value = pStage2LogProbs[i] + static_cast(diversityRate * index); - kvLocal = argmax(kvLocal, {i, value}); - smemVal[i] = value; - } - __syncthreads(); - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage smemReduceBuffer; - __shared__ int threadToUpdate; - - for (int i = 0; i < 2 * nBM; ++i) - { - KVPair kv = BlockReduce(smemReduceBuffer).Reduce(kvLocal, argmax); - if (tid == 0) - { - smemTopKV[i] = kv; - smemVal[kv.key] = -MAX_T_VAL; - threadToUpdate = kv.key % BLOCK_SIZE; - } - __syncthreads(); - // Only one thread needs to update the old partial before the next block reduce. - // No need to do this in the last iteration. - if (tid == threadToUpdate && i < 2 * nBM - 1) - { - kvLocal.key = nCandidate - 1; - kvLocal.value = -MAX_T_VAL; - for (int index = tid; index < nCandidate; index += BLOCK_SIZE) - { - kvLocal = argmax(kvLocal, {index, smemVal[index]}); - } - } - } - } - - if (tid < nBM) // Prepare cumLogProbs for later use - { - smemCumLogProbs[tid] = bh.cumLogProbs[slot * nBM + tid]; - } - __syncthreads(); - - // Timestep at which each next-step destination beam's token is stored in the work tree. - // This is the parent beam's sequence length (the true generation step), which may differ - // from the destination slot's own (possibly stale) length when a finished slot is reused. - __shared__ int smemWriteStep[PBM]; - - if (tid == 0) - { - int nBeamForNextStep{0}; - // Select finished beams into CBA or select tokens for next step sequentially - // Reference (might be changed along HF in the future): - // https://github.com/huggingface/transformers/blob/main/src/transformers/generation/beam_search.py#L272 - for (int i = 0; i < 2 * nBMOut; ++i) - { - int topId; - T topLogProb; - if constexpr (IS_V2) - { - // Get top token and correspongding logProb sequentially from pStage2Ids / pStage2LogProbs - topId = pStage2Ids[i]; - topLogProb = pStage2LogProbs[i]; - } - else - { - // Get top token and correspongding logProb by index of smemTopKV - int const key = smemTopKV[i].key; - topId = pStage2Ids[key]; - topLogProb = pStage2LogProbs[key]; - } - bool const isEndToken = (topId % nV == bh.endIds[slot]); - if (i < nBM && bh.numBeamsCBA != nullptr && isEndToken) - { - // Condition of this branch: - // This token is end-token and belongs to top nBM range in Beam search mode - // Use the actual parent beam index (topId / nV) % nBM, not the candidate rank i, - // to look up the correct sequenceLength and inputLength for length-penalty scoring. - int const parentBeam = (topId / nV) % nBM; - int const nSeqLen - = bh.sequenceLengths[slot * nBM + parentBeam] + 1 - bh.inputLengths[slot * nBM + parentBeam]; - float const score = applyLengthPenalty(topLogProb, nSeqLen, lengthPenalty); - int nCBA = bh.numBeamsCBA[slot]; - if (nCBA >= nBM) - { - // There are already nBM beams - if (score < bh.minNormedScoresCBA[slot]) - { - // Current score is worse than the worst one in candidate beams - if (earlyStopping) - { - // Stop since we have got enough beams - break; - } - else - { - // Continue since there might be longer but better beams - continue; - } - } - else - { - // Current score is better than the worst one in candidate beams - // Find the candidate beam index with the worst score and erase it - for (int j = 0; j < nBM; j++) - { - if (bh.normedScoresCBA[slot * (nBM * 2) + j] == bh.minNormedScoresCBA[slot]) - { - nCBA = j; - bh.numBeamsCBA[slot]--; - bh.minNormedScoresCBA[slot] = FLT_MAX; - bh.normedScoresCBA[slot * (nBM * 2) + j] = score; - for (int l = 0; l < nBM; l++) - { - bh.minNormedScoresCBA[slot] - = min(bh.minNormedScoresCBA[slot], bh.normedScoresCBA[slot * (nBM * 2) + l]); - } - break; - } - } - } - } - // Copy finished beam from work tree to CBA - // The last token - int indexPrev = (topId / nV) % nBM; - int const step = bh.sequenceLengths[slot * nBM + indexPrev]; - int const inputLength = bh.inputLengths[slot * nBM + indexPrev]; - int const offsetCBA = (slot * nBM * 2 + nCBA) * nMSL; - bh.outputIdsCBA[offsetCBA + step] = bh.endIds[slot]; - if (bh.logProbsCBA != nullptr) - { - bh.logProbsCBA[offsetCBA + step] = (float) topLogProb - smemCumLogProbs[(topId / nV) % nBM]; - } - // Previous tokens - for (int j = step - 1; j >= inputLength; j--) - { - bh.outputIdsCBA[offsetCBA + j] = bh.outputIdsPtr[slot][indexPrev * nMSL + j]; - indexPrev = bh.parentIdsPtr[slot][indexPrev * nMSL + j]; - } - if (bh.logProbsCBA != nullptr && bh.logProbsTiled != nullptr) - { - indexPrev = (topId / nV) % nBM; - for (int j = step - 1; j >= inputLength; j--) - { - int const index = (j * nMBS + slot) * nBM + indexPrev; - bh.logProbsCBA[offsetCBA + j] = bh.logProbsTiled[index]; - indexPrev = bh.parentIdsPtr[slot][indexPrev * nMSL + j]; - } - } - // Other parameters - int const index = slot * (nBM * 2) + nCBA; - bh.sequenceLengthsCBA[index] = step; - bh.normedScoresCBA[index] = score; - bh.minNormedScoresCBA[slot] = min(bh.minNormedScoresCBA[slot], bh.normedScoresCBA[index]); - bh.numBeamsCBA[slot]++; - bh.cumLogProbsCBA[index] = (float) topLogProb; - } - else if (i < nBM || bh.numBeamsCBA != nullptr && !isEndToken) - { - // Condition of this branch - // 1. bh.numBeamsCBA == nullptr && i < nBM, i.e., beam search is disable - // 2. bh.numBeamsCBA != nullptr && i < nBM && isEndToken == false, i.e., add token at the end - // 3. bh.numBeamsCBA != nullptr && i >= nBM && isEndToken == false, i.e., add token at the end - // Write at the parent beam's sequence length (the actual generation step), - // not the destination slot's length, which can be stale if the slot was - // previously finished and is now being reused for a new continuation. - int const parentBeam = topId / nV % nBM; - int const step = bh.sequenceLengths[slot * nBM + parentBeam]; - smemWriteStep[nBeamForNextStep] = step; - // Copy the selected token to work tree - bh.outputIdsPtr[slot][nBeamForNextStep * nMSL + step] = topId; - if (bh.logProbsTiled != nullptr) - { - int const index = step * nMBS * nBM + slot * nBM + nBeamForNextStep; - int const indexBeam = parentBeam; - bh.logProbsTiled[index] = (float) topLogProb - smemCumLogProbs[indexBeam]; - } - bh.cumLogProbs[slot * nBM + nBeamForNextStep] = (float) topLogProb; - nBeamForNextStep++; - } - else - { - // Condition of this branch, which we do nothing for it - // 1. bh.numBeamsCBA == nullptr && i >= nBM, i.e., beam search is disable - // 2. bh.numBeamsCBA != nullptr && i >= nBM && isEndToken == true, i.e., ignore the worse beams - } - - if (nBeamForNextStep >= nBMOut) - { - // Condition of this branch - // 1. In EarlyStopping mode, and get enough candidate beams - // 2. In EarlyStopping mode, and get enough tokens for the next generation step - // 3. In NonEarlyStopping mode, and get enough tokens for the next generation step - // TODO: improve the condition like below - // earlyStopping == 1 && bh.numBeamsCBA[slot] >= nBM || nBeamForNextStep >= nBM - break; - } - } - smemNBeamForNextStep = nBeamForNextStep; - } - - // Update bh.batchDones - if (tid == 0 && bh.numBeamsCBA != nullptr) - { - if (bh.numBeamsCBA[slot] < nBM) - { - // no enough beams - bh.batchDones[slot] = false; - } - else if (earlyStopping == 1) - { - // enough candidate beams in EarlyStopping mode - bh.batchDones[slot] = true; - } - else - { - // enough beams in NonEarlyStopping mode - int nSeqLen = bh.sequenceLengths[slot * nBM] + 1 - bh.inputLengths[slot * nBM]; - float const bestCumLogProbs = (IS_V2) ? pStage2LogProbs[0] : smemTopKV[0].value; - // According to semantics of HF, smemTopKV[0].value is used as bestCumLogProbs - // But maybe bh.cumLogProbs[slot * nBM + i] is more suitable? - // https://github.com/huggingface/transformers/blob/main/src/transformers/generation/beam_search.py#L307 - if (earlyStopping != 0 && lengthPenalty > 0.0f) - { - // Specialization for earlyStopping == "never" and lengthPenalty > 0 in HF - nSeqLen = nMSL - bh.inputLengths[slot * nBM]; - } - float const bestAttainableScore = applyLengthPenalty(bestCumLogProbs, nSeqLen, lengthPenalty); - bh.batchDones[slot] = bh.minNormedScoresCBA[slot] >= bestAttainableScore; - } - } - __syncthreads(); - - // Update sequenceLengths, parentIdsPtr, outputIdsPtr and finished - if (tid < nBM) - { - smemSeqLen[tid] = bh.sequenceLengths[slot * nBM + tid]; - } - __syncthreads(); - - if (tid < nBMOut) - { - int const indexBatchBeam = slot * nBM + tid; - if (tid < smemNBeamForNextStep) - { - // This slot received a valid next-step token from the selection phase. - // Use the timestep recorded by the selection phase (the parent beam's length), - // which matches the position where the encoded token was stored. - int const step = smemWriteStep[tid]; - int const newId = bh.outputIdsPtr[slot][tid * nMSL + step]; - int const newBeamId = (newId / nV) % nBM; - int const newTokenId = newId % nV; - int const indexParentBeam = slot * nBM + newBeamId; - int const parentSeqLen = smemSeqLen[newBeamId]; - bh.sequenceLengths[indexBatchBeam] = parentSeqLen + (!bh.finished[indexParentBeam].isFinished() ? 1 : 0); - if (newTokenId == bh.endIds[slot]) - { - bh.finished[indexBatchBeam].setFinishedEOS(); - } - else - { - // Reset any stale finished state: this slot may have been marked finished in a - // previous step and is now reused for a valid non-EOS beam; otherwise it would be - // wrongly skipped downstream. - bh.finished[indexBatchBeam] = FinishedState::empty(); - } - bh.parentIdsPtr[slot][tid * nMSL + step] = newBeamId; - bh.outputIdsPtr[slot][tid * nMSL + step] = newTokenId; - } - else - { - // No valid next-step token for this slot: all top candidates went to CBA. - // Mark as finished so downstream stages (cache indirection, next decode) skip it. - bh.finished[indexBatchBeam].setFinished(); - } - - if ((earlyStopping == 1) && (bh.numBeamsCBA != nullptr && bh.numBeamsCBA[slot] >= nBM) - || (earlyStopping != 1) && bh.batchDones[slot]) - { - bh.batchDones[slot] = true; - bh.finished[indexBatchBeam].setFinished(); - } - } -} - -template -void beamSearchKernelLauncher( - T const* logProbs, T const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream) -{ - // clang-format off - /* - V1 Workflow (reference: https://github.com/NVIDIA/online-softmax): - logProbs.shape = [nBS, nBM, nV] - nV |<- nVChunk ->|<- nVChunk ->| <- ... ->| |<- nBM*4 ->|<- nBM*4 ->|<- ... ->| ■ - ┏━━━━━━━━━━┓ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ - ┃nBM ┃ ┃nBM ┃ ┃nBM ┃ - ┣━━━━━━━━━━┫ ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫ A ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫ B - nBS ┃nBM ┃ ---> nBS ┃nBM ┃ ---> nBS ┃nBM ┃ ---> - ┣━━━━━━━━━━┫ ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫ ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫ - ┃nBM ┃ ┃nBM ┃ ┃nBM ┃ - ┗━━━━━━━━━━┛ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ - logProbs divide `nV` elements into `nVPart` parts pStage3 with `nVPart` tiles per row - - |<- nBm*2 ->| |<- nBm*2 ->| - ┏━━━━━━━━━━━┓ ┏━━━━━━━━━━━┓ - ┃nBM ┃ ┃nBM ┃ - B ┣━━━━━━━━━━━┫ ┣━━━━━━━━━━━┫ C - ---> nBS ┃nBM ┃ ┃nBM ┃ ---> - ┣━━━━━━━━━━━┫ ┣━━━━━━━━━━━┫ - ┃nBM ┃ ┃nBM ┃ - ┗━━━━━━━━━━━┛ ┗━━━━━━━━━━━┛ - pStage2Ids pStage2LogProbs - - ■: Each "tile" in pStage3 with shape [`nBM*4`] contains `nBM*2` top ids and corresponding `nBM*2` log probs. - |<- nBm*2 ->|<- nBm*2 ->| - ┏━━━━━━━━━━━━━━━━━━━━━━━┓ - 1 ┃ top ids | log probs ┃ - ┗━━━━━━━━━━━━━━━━━━━━━━━┛ - - A: beamStage1Kernel: gridDim(BS,BM,nVPart), blockDim(nThreadStage1,1,1) - Each Block takes `nVChunk` contiguous elements from `logProbs`, does TopK and writes output to `pStage3` - B: beamStage2Kernel: gridDim(BS,BM,1), blockDim(32/64/128,1,1) - Each Block takes `nVPart` contiguous tiles from pStage3, add `cumLogProbs`, does TopK` and writes output to `pStage2Ids` and `pStage2LogProbs` - C: beamStage3Kernel: gridDim(BS,1,1), blockDim(128,1,1) - Main logic of Beam-Search, each Block is responsible for one batch, doing work below: - + moves one beam into candidate-beam-array if it is finished (gemerated end_id in this step). - + selects BM elements for the next generation step if not. - + maintains related score array, min_normed_score / batchDones / finished, etc.. - - =================================================================================================================================== - - V2 Workflow (use Air-TopK for better performance, https://dl.acm.org/doi/pdf/10.1145/3581784.3607062) - logProbs.shape = [nBS, nBM, nV] - |<- nV ->| |<- nBM*2 ->| |<- nBM*2 ->| |<- nBM*2 ->| |<- nBM*2 ->| |<- nBM*2 ->| - ┏━━━━━━━━┓ ┏━━━━━━━━━━━┓ ┏━━━━━━━━━━━┓ ┏━━━━━━━━━━━┓ ┏━━━━━━━━━━━┓ D ┏━━━━━━━━━━━┓ - ┃nBM ┃ ┃nBM ┃ ┃nBM ┃ ┃nBM ┃ nBS ┃ ┃ ---> nBS ┃ ┃ ---\ - ┣━━━━━━━━┫ A ┣━━━━━━━━━━━┫ ┣━━━━━━━━━━━┫ B ┣━━━━━━━━━━━┫ C ┗━━━━━━━━━━━┛ ┗━━━━━━━━━━━┛ | E - nBS ┃nBM ┃ ---> nBS ┃nBM ┃ ┃nBM ┃ ---> nBS ┃nBM ┃ ---> pStage2Id pStage2Id |---> - ┣━━━━━━━━┫ ┣━━━━━━━━━━━┫ ┣━━━━━━━━━━━┫ ┣━━━━━━━━━━━┫ ┏━━━━━━━━━━━┓ | - ┃nBM ┃ ┃nBM ┃ ┃nBM ┃ ┃nBM ┃ nBS ┃ ┃ --------------------------/ - ┗━━━━━━━━┛ ┗━━━━━━━━━━━┛ ┗━━━━━━━━━━━┛ ┗━━━━━━━━━━━┛ ┗━━━━━━━━━━━┛ - logProbs pStage1Id pStage1LogProbs pStage1LogProbs pStage2LogProbs - - A: TopK : Get top `nBM*2` elements in `nBS*nBM` groups (`nV` elements per group) - B: addCumLogProbs : Add `cumLogProbs` to the elements in each beam - C: TopK : Get top `nBM*2` elements in `nBS` group (`nBM*nBM*2` elements per group) - D: gatherIds : Combine stage1Id and stage2Id to get ids of the top `nBM*2` elements in input logProbs - E: beamStage3Kernel: Main logic of Beam-Search, each Block is responsible for one batch, doing work below: - + moves one beam into candidate-beam-array if it is finished (gemerated end_id in this step). - + selects BM elements for the next generation step if not. - + maintains related score array, min_normed_score / batchDones / finished, etc.. - - =================================================================================================================================== - - V2 Workflow for VBWS, similar to V2 workflow above, but `nBMIn` and `nBMOut` might be different from `nBM` - logProbs.shape = [nBS, nBMIn, nV] - |<- nV ->| |<- nBMOut*2 ->| |<- nBMOut*2 ->| |<- nBMOut*2 ->| |<- nBMOut*2 ->| |<- nBMOut*2 ->| - ┏━━━━━━━━┓ ┏━━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━━━┓ D ┏━━━━━━━━━━━━━━┓ - ┃nBMIn ┃ ┃nBMIn ┃ ┃nBMIn ┃ ┃nBMIn ┃ nBS ┃ ┃ ---> nBS ┃ ┃ ---\ - ┣━━━━━━━━┫ A ┣━━━━━━━━━━━━━━┫ ┣━━━━━━━━━━━━━━┫ B ┣━━━━━━━━━━━━━━┫ C ┗━━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━━━┛ | E - nBS ┃nBMIn ┃ ---> nBS ┃nBMIn ┃ ┃nBMIn ┃ ---> nBS ┃nBMIn ┃ ---> pStage2Id pStage2Id |---> - ┣━━━━━━━━┫ ┣━━━━━━━━━━━━━━┫ ┣━━━━━━━━━━━━━━┫ ┣━━━━━━━━━━━━━━┫ ┏━━━━━━━━━━━━━━┓ | - ┃nBMIn ┃ ┃nBMIn ┃ ┃nBMIn ┃ ┃nBMIn ┃ nBS ┃ ┃ -----------------------------/ - ┗━━━━━━━━┛ ┗━━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━━━┛ - logProbs pStage1Id pStage1LogProbs pStage1LogProbs pStage2LogProbs - */ - // clang-format on - - size_t const nBS{bh.nBatchSize}; - size_t const nBM{bh.nBeamWidth}; - size_t const nV{bh.nVocabSize}; - size_t const nVPart{bh.nVPart}; - size_t const nByteMaxSharedMemoryPerBlock{bh.nByteMaxSharedMemoryPerBlock}; - int* pStage2Ids{nullptr}; - T* pStage2LogProbs{nullptr}; - float* pStage3{nullptr}; - - // VBWS: - // + `nBMIn` / `nBMOut` is the beam width in the last / next network forward computation respectively - // + `nBM` is the max value of the beam width array, which is used for memory allocatation - // Normal Beam Search: - // + `nBMIn` / `nBMOut` / `nBM` share the same value - // TODO: now `nBMIn` and `nBMOut` of request 0 is used for the whole batch, - // change to corresponding BMs if Diverse-Beam-Width-Search is supported - size_t const nBMIn = bh.bVBWS ? bh.nBeamWidthInHost[0] : nBM; - size_t const nBMOut = bh.bVBWS ? bh.nBeamWidthOutHost[0] : nBM; - bh.nBeamWidthIn = nBMIn; // Save nBMIn back to bh - bh.nBeamWidthOut = nBMOut; // Save nBMOut back to bh - - if constexpr (IS_V2) - { - // see `BeamSearchLayer::configureBeamSearchLayer()` for the workspace structure - size_t const offsetStage1 = roundUp(nBS * nBM * nBM * 2, 4); - size_t const offsetStage2 = roundUp(nBS * nBM * 2, 4); - pStage2Ids = reinterpret_cast(workspace); - int offset = sizeof(int) * offsetStage2; - pStage2LogProbs = reinterpret_cast(reinterpret_cast(workspace) + offset); - offset += sizeof(T) * offsetStage2; - int* pStage1Ids = reinterpret_cast(reinterpret_cast(workspace) + offset); - pStage3 = reinterpret_cast(reinterpret_cast(workspace) + offset); - offset += sizeof(int) * offsetStage1; - T* pStage1LogProbs = reinterpret_cast(reinterpret_cast(workspace) + offset); - offset += sizeof(T) * offsetStage1; - void* pTopK = reinterpret_cast(reinterpret_cast(workspace) + offset); - - // Stage 1 - invokeTopkLastDim(nBS * nBMIn, nV, nBMOut * 2, true, logProbs, pStage1LogProbs, pStage1Ids, pTopK, stream); - sync_check_cuda_error(stream); - - int nThread = min(roundUp(nBMIn * nBMOut * 2, 32), 1024); - addCumLogProbs<<>>(pStage1LogProbs, pStage1Ids, bh.cumLogProbs, bh.finished, bh.endIds, - bh.diversityRates, bh.batchSlots, nBS, nBMIn, nBMOut, nBM); - sync_check_cuda_error(stream); - - // Stage 2 - invokeTopkLastDim( - nBS, nBMIn * nBMOut * 2, nBMOut * 2, true, pStage1LogProbs, pStage2LogProbs, pStage2Ids, pTopK, stream); - sync_check_cuda_error(stream); - - nThread = min(roundUp(nBMOut * 2, 32), 1024); - gatherId<<>>(pStage1Ids, pStage2Ids, nBS, nBMIn, nBMOut, nV); - sync_check_cuda_error(stream); - } - else // V1 - { - // see `BeamSearchLayer::configureBeamSearchLayer()` for the workspace structure - int const offset = roundUp(nBS * nBM * nBM * 2, 4); - pStage2Ids = reinterpret_cast(workspace); - pStage2LogProbs = reinterpret_cast(pStage2Ids + offset); - pStage3 = reinterpret_cast(pStage2LogProbs + offset); - - // Stage 1 - size_t constexpr nThreadStage1 = (PBM < 16) ? ((PBM < 8) ? kThreadForSmallBeamWidth : 128) : 64; - dim3 grid(nBS, nBM, bh.nVPart), block(nThreadStage1); - beamStage1Kernel<<>>( - logProbs, bias, pStage3, bh.endIds, bh.finished, nV, bh.batchSlots); - sync_check_cuda_error(stream); - -// Stage 2 -#define BEAM_STAGE2_KERNEL(N_VOCAB_PART, IS_FAST) \ - { \ - if (IS_FAST && nByteRuntimeSharedMemory > (48 << 10)) \ - { \ - TLLM_CUDA_CHECK(cudaFuncSetAttribute(beamStage2Kernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, nByteRuntimeSharedMemory)); \ - } \ - beamStage2Kernel \ - <<>>( \ - pStage2Ids, pStage2LogProbs, pStage3, bh.cumLogProbs, bh.batchSlots, nV, nVPart); \ - } - // TODO: rewrite kernel to remove dependence of constant block size to reduce compilation time - size_t nByteRuntimeSharedMemory - = sizeof(float) * nVPart * (PBM * 4) + sizeof(cub::KeyValuePair) * PBM * 2; - if (nByteRuntimeSharedMemory <= nByteMaxSharedMemoryPerBlock && nVPart <= 32) - { - BEAM_STAGE2_KERNEL(32, true) - } - else if (nByteRuntimeSharedMemory <= nByteMaxSharedMemoryPerBlock && nVPart <= 64) - { - BEAM_STAGE2_KERNEL(64, true) - } - else if (nByteRuntimeSharedMemory <= nByteMaxSharedMemoryPerBlock) - { - BEAM_STAGE2_KERNEL(128, true) - // No branch with larger `N_VOCAB_PART` since nVPart <= kMaxVPartStage1 == 128 - } - else - { - TLLM_LOG_TRACE("Use slow Beam Search stage 2 kernel due to large beam_width or vocab_size"); - BEAM_STAGE2_KERNEL(128, false) - } - sync_check_cuda_error(stream); -#undef BEAM_STAGE2_KERNEL - } - - // Stage 3 in common - size_t constexpr nThreadStage3 = (PBM + 31) / 32 * 32; - size_t const nByteStaticSharedMemory = bh.nByteSharedMemoryStage3; - size_t const nByteDynamicSharedMemory = (IS_V2) ? 0 : sizeof(T) * nBM * nBM * 2; - size_t const nByteRuntimeSharedMemory = nByteStaticSharedMemory + nByteDynamicSharedMemory; - - if (nByteRuntimeSharedMemory <= nByteMaxSharedMemoryPerBlock) - { - if (nByteRuntimeSharedMemory > (48 << 10)) - { - TLLM_CUDA_CHECK(cudaFuncSetAttribute(beamStage3Kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, nByteRuntimeSharedMemory)); - } - beamStage3Kernel - <<>>(pStage2Ids, pStage2LogProbs, pStage3, bh); - } - else - { - if (nByteStaticSharedMemory > (48 << 10)) - { - TLLM_CUDA_CHECK(cudaFuncSetAttribute(beamStage3Kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, nByteStaticSharedMemory)); - } - beamStage3Kernel - <<>>(pStage2Ids, pStage2LogProbs, pStage3, bh); - } - sync_check_cuda_error(stream); - - return; -} - -#undef BEAM_STAGE2_KERNEL - -#define INSTANTIATE_BEAM_SEARCH(T, PBM, IS_V2) \ - template void beamSearchKernelLauncher( \ - T const* logProbs, T const* bias, void* workspace, BeamHypotheses& bh, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/decodingKernels.cu b/cpp/tensorrt_llm/kernels/decodingKernels.cu deleted file mode 100644 index 7c72091abad2..000000000000 --- a/cpp/tensorrt_llm/kernels/decodingKernels.cu +++ /dev/null @@ -1,835 +0,0 @@ -/* - * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" -#include "tensorrt_llm/kernels/decodingKernels.h" - -#ifndef CUDART_VERSION -#error CUDART_VERSION Undefined! -#elif (CUDART_VERSION >= 11050) -#include -#else -#include "3rdparty/cub/cub.cuh" -#endif - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::runtime; - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -class CopyBeamHypothesesStruct -{ -public: - TokenIdType const* srcOutputIdsCBA; // [BS, BM*2, MSL] - TokenIdType* dstOutputIdsCBA; // [BS, BM*2, MSL] - SizeType32 outputIdsNumElts; - - float const* srcLogProbsCBA; // [BS, BM*2, MSL] - float* dstLogProbsCBA; // [BS, BM*2, MSL] - SizeType32 logProbsNumElts; - - SizeType32 const* srcSequenceLengthsCBA; // [BS, BM*2] - SizeType32* dstSequenceLengthsCBA; // [BS, BM*2] - SizeType32 sequenceLengthsNumElts; - - float const* srcCumLogProbsCBA; // [BS, BM*2] - float* dstCumLogProbsCBA; // [BS, BM*2] - SizeType32 cumLogProbsCBANumElts; - - float const* srcNormedScoresCBA; // [BS, BM*2] - float* dstNormedScoresCBA; // [BS, BM*2] - SizeType32 normedScoresNumElts; - - SizeType32 const* srcNumBeamsCBA; // [BS] - SizeType32* dstNumBeamsCBA; // [BS] - SizeType32 numBeamsNumElts; - - float const* srcMinNormedScoresCBA; // [BS] - float* dstMinNormedScoresCBA; // [BS] - SizeType32 minNormedScoresNumElts; - - bool const* srcBatchDones; // [BS] - bool* dstBatchDones; // [BS] - SizeType32 batchDonesNumElts; - - float const* srcCumLogProbs; // [BS, BM] - float* dstCumLogProbs; // [BS, BM] - SizeType32 cumLogProbsNumElts; -}; - -__global__ void gatherTree(gatherTreeParam param) -{ - for (int batchbeamIdx = blockIdx.x * blockDim.x + threadIdx.x; batchbeamIdx < param.batchSize * param.beamWidth; - batchbeamIdx += gridDim.x * blockDim.x) - { - int const batch = batchbeamIdx / param.beamWidth; - int const beam = batchbeamIdx % param.beamWidth; - int const inputLen = param.inputLengths == nullptr ? 0 : param.inputLengths[batchbeamIdx]; - - int const* parentIds = param.parentIds; - int const* stepIds = param.stepIds; - - // TODO optimize the reduce_max operation for large beamWidth - int maxLen = -1; - bool updateResponseInputLength = param.responseInputLengths != nullptr; - // int selected_beam_index = 0; - for (int beamIdx = 0; beamIdx < param.beamWidth; beamIdx++) - { - int tmpLen - = param.sequenceLengths[batch * param.beamWidth + beamIdx] + param.maxSequenceLengthFinalStep - 1; - param.sequenceLengths[batch * param.beamWidth + beamIdx] = tmpLen; - if (updateResponseInputLength) - { - param.responseInputLengths[batch * param.beamWidth + beamIdx] = inputLen; - } - if (tmpLen > maxLen) - { - maxLen = tmpLen; - } - } - int const maxSeqLenB = min(param.maxSeqLen, maxLen); - if (maxSeqLenB <= 0) - { - continue; - } - - int const initialTgtIx = batch * param.beamWidth * param.maxSeqLen + beam * param.maxSeqLen + maxSeqLenB - 1; - int const initialParentIx = batch * param.beamWidth * param.maxSeqLen + beam * param.maxSeqLen + maxSeqLenB - 1; - param.outputIds[initialTgtIx] = __ldg(stepIds + initialParentIx); - int parent = parentIds == nullptr ? 0 : __ldg(parentIds + initialParentIx) % param.beamWidth; - bool foundBad = false; - - for (int level = maxSeqLenB - 2; level >= 0; --level) - { - int const levelBeamIx = batch * param.beamWidth * param.maxSeqLen + beam * param.maxSeqLen + level; - int const levelParentIx = batch * param.beamWidth * param.maxSeqLen + parent * param.maxSeqLen + level; - if (parent < 0 || parent >= param.beamWidth) - { - param.outputIds[levelBeamIx] = param.endTokens[batch]; - parent = -1; - foundBad = true; - } - else - { - param.outputIds[levelBeamIx] = __ldg(stepIds + levelParentIx); - parent = parentIds == nullptr ? 0 : __ldg(parentIds + levelParentIx) % param.beamWidth; - } - } - // set the padded part as end_token - // inputLen - for (int index = maxLen; index < param.maxSeqLen; ++index) - { - param.outputIds[batch * param.beamWidth * param.maxSeqLen + beam * param.maxSeqLen + index] - = param.endTokens[batch]; - } - - // Not necessary when using a BeamSearchDecoder, but necessary - // when a user feeds in possibly broken trajectory (i.e., non-eos - // entries in a beam following eos entries). - if (!foundBad) - { - bool finished = false; - // skip the step 0 because it is often the start token - int startStep = 1; - for (int time = startStep; time < maxSeqLenB; ++time) - { - int const levelBeamIx = batch * param.beamWidth * param.maxSeqLen + beam * param.maxSeqLen + time; - if (finished) - { - param.outputIds[levelBeamIx] = param.endTokens[batch]; - } - else if (param.outputIds[levelBeamIx] == param.endTokens[batch]) - { - finished = true; - } - } - } - } -} - -struct RankNorm -{ - int rank; - float norm; -}; - -inline __device__ RankNorm swap(RankNorm const& rankNorm, int mask, int dir) -{ - // Exchange RankNorm data inside the warp - RankNorm other; - other.rank = __shfl_xor_sync(unsigned(-1), rankNorm.rank, mask); - other.norm = __shfl_xor_sync(unsigned(-1), rankNorm.norm, mask); - // dir == 0 -> return larger one - // dir == 1 -> return smaller one - bool doSwap = (rankNorm.norm != other.norm) && ((rankNorm.norm > other.norm) == dir); - return doSwap ? other : rankNorm; -} - -inline __device__ uint32_t bfe(uint32_t a, uint32_t start, uint32_t len = 1) -{ - uint32_t d; - asm volatile("bfe.u32 %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(start), "r"(len)); - return d; -} - -__global__ void finalized(gatherTreeParam param) -{ - int const beamIdx = static_cast(threadIdx.x); - int const beamWidth{param.beamWidth}; - - extern __shared__ char array[]; - int* sRank = (int*) (array); - int* sLength = (int*) (sRank + beamWidth); - float* sScores = (float*) (sLength + beamWidth); - float* sNormedScores = (float*) (sScores + beamWidth); - int* sIds = (int*) (sNormedScores + beamWidth); - - if (beamIdx < beamWidth) - { - int const idx = blockIdx.x * param.beamWidth + beamIdx; - int const numGeneratedToken{param.sequenceLengths[idx] - param.inputLengths[idx]}; - sNormedScores[beamIdx] = applyLengthPenalty(param.cumLogProbs[idx], numGeneratedToken, param.lengthPenalty); - sLength[beamIdx] = param.sequenceLengths[idx]; - sScores[beamIdx] = param.cumLogProbs[idx]; - } - for (int idx = beamIdx; idx < beamWidth * param.maxSeqLen; idx += blockDim.x) - { - sIds[idx] = param.outputIds[blockIdx.x * param.beamWidth * param.maxSeqLen + idx]; - } - __syncthreads(); - - RankNorm rankNorm; - rankNorm.rank = beamIdx; - rankNorm.norm = beamIdx < beamWidth ? sNormedScores[beamIdx] : -FLT_MAX; - - if (beamWidth < 32) - { - int warpid = threadIdx.x / 32; - int laneid = threadIdx.x % 32; - - if (warpid == 0 && beamWidth > 1) - { - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 1) ^ bfe(laneid, 0)); // 2 - } - - if (warpid == 0 && beamWidth > 2) - { - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 2) ^ bfe(laneid, 1)); // 3~4 - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 2) ^ bfe(laneid, 0)); - } - - if (warpid == 0 && beamWidth > 4) - { - rankNorm = swap(rankNorm, 0x04, bfe(laneid, 3) ^ bfe(laneid, 2)); // 5~8 - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 3) ^ bfe(laneid, 1)); - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 3) ^ bfe(laneid, 0)); - } - - if (warpid == 0 && beamWidth > 8) - { - rankNorm = swap(rankNorm, 0x08, bfe(laneid, 4) ^ bfe(laneid, 3)); // 9~16 - rankNorm = swap(rankNorm, 0x04, bfe(laneid, 4) ^ bfe(laneid, 2)); - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 4) ^ bfe(laneid, 1)); - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 4) ^ bfe(laneid, 0)); - } - - if (warpid == 0 && beamWidth > 16) - { - rankNorm = swap(rankNorm, 0x10, bfe(laneid, 5) ^ bfe(laneid, 4)); // 17~32 - rankNorm = swap(rankNorm, 0x08, bfe(laneid, 5) ^ bfe(laneid, 3)); - rankNorm = swap(rankNorm, 0x04, bfe(laneid, 5) ^ bfe(laneid, 2)); - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 5) ^ bfe(laneid, 1)); - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 5) ^ bfe(laneid, 0)); - } - } - else - { - // Not supported! We must have a check before calling that kernel. - } - - if (beamIdx < beamWidth) - { - sRank[beamIdx] = rankNorm.rank; - } - - __syncthreads(); - - if (beamIdx < beamWidth) - { - auto srcIdx{rankNorm.rank}; - auto tgtIdx{blockIdx.x * param.beamWidth + beamIdx}; - param.sequenceLengths[tgtIdx] = sLength[srcIdx]; - param.cumLogProbs[tgtIdx] = sScores[srcIdx]; - } - - for (int beamIdx = 0; beamIdx < beamWidth; beamIdx++) - { - for (int i = threadIdx.x; i < sLength[sRank[beamIdx]]; i += blockDim.x) - { - param.outputIds[blockIdx.x * beamWidth * param.maxSeqLen + beamIdx * param.maxSeqLen + i] - = sIds[sRank[beamIdx] * param.maxSeqLen + i]; - } - } -} - -void invokeGatherTree(gatherTreeParam param) -{ - int batchbeam = param.batchSize * param.beamWidth; - dim3 grid(1), block(batchbeam); - // though decoder do not support > 1024 for now - if (batchbeam > 1024) - { - grid.x = ceil(param.batchSize * param.beamWidth / 1024.); - block.x = 1024; - } - gatherTree<<>>(param); - sync_check_cuda_error(param.stream); - - if (param.beamWidth > 1) - { - TLLM_CHECK_WITH_INFO(param.beamWidth <= 32, "TRT-LLM does not support beam width > 32 now"); - // sort results by normalized cumLogProbs - dim3 grid(param.batchSize); - dim3 block(divUp(param.beamWidth, 32) * 32); - - auto shm_size = param.beamWidth * (sizeof(float) * 2 + sizeof(int) * 2 + sizeof(int) * param.maxSeqLen); - finalized<<>>(param); - } -} - -__global__ void insertUnfinishedPathKernel(BeamHypotheses bh) -{ - // Move ALL unfinished beams from bh.outputIdsUnfinish to bh.outputIdsCBA - // So here might be more than `nBM` beams in bh.outputIdsCBA after this kernel - // Data movement: - // bh.outputIdsUnfinish -> bh.outputIdsCBA - // bh.sequenceLengths -> bh.sequenceLengthsCBA - // bh.cumLogProbs -> bh.cumLogProbsCBA - // bh.logProbsTiled -> bh.logProbsCBA - // update bh.normedScoresCBA - // update bh.numBeamsCBA - - size_t const bid = blockIdx.x; // Index of Batch - size_t const nBM{bh.nBeamWidth}; - size_t const nMBS{bh.nMaxBatchSize}; // Only for bh.logProbsTiled - size_t const nMSL{bh.nMaxSeqLen}; - bool const bOutputLogProbs{bh.logProbsCBA != nullptr && bh.logProbsTiled != nullptr}; - int const indexDstStart{bh.numBeamsCBA[bid]}; - - if (bh.batchDones[bid]) - { - return; - } - - for (int i = 0; i < nBM; ++i) - { - int const srcBeam = bid * nBM + i; - int const dstBeam = bid * nBM * 2 + i + indexDstStart; - int const step = bh.sequenceLengths[srcBeam] - 1; - int const inputLength = bh.inputLengths[srcBeam]; - - // The last token - int const srcId = srcBeam * nMSL + step; - int const dstId = dstBeam * nMSL + step; - bh.outputIdsCBA[dstId] = bh.outputIdsUnfinish[srcId]; - if (bOutputLogProbs) - { - bh.logProbsCBA[dstId] = bh.logProbsTiled[step * nMBS * nBM + srcBeam]; - } - // Previous tokens - int prevId = bh.parentIdsUnfinish[srcId]; - for (int j = step - 1; j >= inputLength; --j) - { - int const index = bid * nBM * nMSL + prevId * nMSL + j; - bh.outputIdsCBA[dstBeam * nMSL + j] = bh.outputIdsUnfinish[index]; - prevId = bh.parentIdsUnfinish[index]; - } - if (bOutputLogProbs) - { - prevId = bh.parentIdsUnfinish[srcId]; - for (int j = step - 1; j >= inputLength; --j) - { - int const index = bid * nBM * nMSL + prevId * nMSL + j; - bh.logProbsCBA[dstBeam * nMSL + j] = bh.logProbsTiled[j * nMBS * nBM + bid * nBM + prevId]; - prevId = bh.parentIdsUnfinish[index]; - } - } - // Other parameters - bh.sequenceLengthsCBA[dstBeam] = bh.sequenceLengths[srcBeam]; - bh.normedScoresCBA[dstBeam] - = applyLengthPenalty(bh.cumLogProbs[srcBeam], step - bh.inputLengths[srcBeam] + 1, bh.lengthPenalties[bid]); - bh.cumLogProbsCBA[dstBeam] = bh.cumLogProbs[srcBeam]; - bh.numBeamsCBA[bid]++; - } -} - -void invokeInsertUnfinishedPath(BeamHypotheses& bh, cudaStream_t stream) -{ - insertUnfinishedPathKernel<<>>(bh); -} - -__global__ void finalizeKernel(BeamHypotheses bh) -{ - // Do index sort on bh.normedScoresCBA, then move buffers from CBA to output by the order of index - // Data movement: - // bh.outputIdsCBA -> bh.outputIds - // bh.sequenceLengthsCBA -> bh.sequenceLengths - // bh.cumLogProbsCBA -> bh.cumLogProbs - // bh.logProbsCBA -> bh.logProbs - - int const bid = blockIdx.x; // Index of Batch - int const tid = threadIdx.x; // Index of Beam - size_t const nBM{bh.nBeamWidth}; - size_t const nMSL{bh.nMaxSeqLen}; - int const nCBA{bh.numBeamsCBA[bid]}; // Count of candidates in CBA, nBM <= nCBA <= 2*nBM - - extern __shared__ char smem[]; - int* smemRank = (int*) (smem); // [nBM] - float* smemScore = (float*) (smemRank + nBM); // [2*nBM] - int* smemSL = (int*) (smemScore + nBM * 2); // [nBM] - - // Sort - for (int i = tid; i < nCBA; i += blockDim.x) - { - smemScore[i] = bh.normedScoresCBA[bid * nBM * 2 + i]; - } - __syncthreads(); - - if (nCBA <= 32) - { - int const warpid = tid / 32; - int const laneid = tid % 32; - RankNorm rankNorm{tid, tid < nCBA ? smemScore[tid] : -FLT_MAX}; - - if (warpid == 0 && nCBA > 1) - { - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 1) ^ bfe(laneid, 0)); // 2 - } - if (warpid == 0 && nCBA > 2) - { - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 2) ^ bfe(laneid, 1)); // 3~4 - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 2) ^ bfe(laneid, 0)); - } - if (warpid == 0 && nCBA > 4) - { - rankNorm = swap(rankNorm, 0x04, bfe(laneid, 3) ^ bfe(laneid, 2)); // 5~8 - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 3) ^ bfe(laneid, 1)); - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 3) ^ bfe(laneid, 0)); - } - if (warpid == 0 && nCBA > 8) - { - rankNorm = swap(rankNorm, 0x08, bfe(laneid, 4) ^ bfe(laneid, 3)); // 9~16 - rankNorm = swap(rankNorm, 0x04, bfe(laneid, 4) ^ bfe(laneid, 2)); - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 4) ^ bfe(laneid, 1)); - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 4) ^ bfe(laneid, 0)); - } - if (warpid == 0 && nCBA > 16) - { - rankNorm = swap(rankNorm, 0x10, bfe(laneid, 5) ^ bfe(laneid, 4)); // 17~32 - rankNorm = swap(rankNorm, 0x08, bfe(laneid, 5) ^ bfe(laneid, 3)); - rankNorm = swap(rankNorm, 0x04, bfe(laneid, 5) ^ bfe(laneid, 2)); - rankNorm = swap(rankNorm, 0x02, bfe(laneid, 5) ^ bfe(laneid, 1)); - rankNorm = swap(rankNorm, 0x01, bfe(laneid, 5) ^ bfe(laneid, 0)); - } - if (tid < nBM) - { - smemRank[tid] = rankNorm.rank; - } - __syncthreads(); - } - else - { - for (int i = 0; i < nBM; ++i) - { - float maxScore = -FLT_MAX; - for (int j = 0; j < (nCBA + 1024 - 1) / 1024; ++j) - { - int const index = tid + 1024 * j; - float const score = (index < bh.numBeamsCBA[bid]) ? smemScore[index] : -FLT_MAX; - float const maxScore1 = blockReduceMax(score); - maxScore = max(maxScore, maxScore1); - } - if (tid == 0) - { - for (int j = 0; j < nCBA; ++j) - { - if (smemScore[j] == maxScore) - { - smemRank[i] = j; - smemScore[j] = -FLT_MAX; - break; - } - } - } - __syncthreads(); - } - } - - // Move bh.sequenceLengths, bh.cumLogProbs - if (tid < nBM) - { - smemSL[tid] = bh.sequenceLengthsCBA[bid * nBM * 2 + smemRank[tid]]; - bh.sequenceLengths[bid * nBM + tid] = smemSL[tid]; - if (bh.cumLogProbs != nullptr) - { - bh.cumLogProbs[bid * nBM + tid] = bh.cumLogProbsCBA[bid * nBM * 2 + smemRank[tid]]; - } - } - __syncthreads(); - - // Move bh.outputIds, bh.logProbs - for (int beamIdx = 0; beamIdx < nBM; beamIdx++) - { - int const inputLength = bh.inputLengths[bid * nBM + beamIdx]; - for (int i = tid; i < smemSL[beamIdx]; i += blockDim.x) - { - int const dst = bid * nBM * nMSL + beamIdx * nMSL + i; - if (i < inputLength) - { - int const src = bid * nBM * nMSL + beamIdx * nMSL + i; - bh.outputIds[dst] = bh.outputIdsUnfinish[src]; - } - else - { - int const src = bid * nBM * 2 * nMSL + smemRank[beamIdx] * nMSL + i; - bh.outputIds[dst] = bh.outputIdsCBA[src]; - } - } - if (bh.logProbs != nullptr) - { - for (int i = tid; i < smemSL[beamIdx]; i += blockDim.x) - { - if (i >= inputLength) - { - int const dst = bid * nBM * nMSL + beamIdx * nMSL + i; - int const src = bid * nBM * 2 * nMSL + smemRank[beamIdx] * nMSL + i; - bh.logProbs[dst - inputLength] = bh.logProbsCBA[src]; - } - } - } - } -} - -void invokeFinalize(BeamHypotheses& bh, cudaStream_t stream) -{ - int const nBM = bh.nBeamWidth; - int const nThread = min(roundUp(nBM * 2, 32), 1024); - size_t const nByteSharedMemory = (sizeof(int) + sizeof(float)) * nBM * 2; - finalizeKernel<<>>(bh); - sync_check_cuda_error(stream); -} - -__global__ void copyBeamHypotheses(CopyBeamHypothesesStruct copyStruct) -{ - auto const idx = static_cast(threadIdx.x + blockIdx.x * blockDim.x); - auto const stride = static_cast(blockDim.x * gridDim.x); - - for (SizeType32 ii = idx; ii < copyStruct.outputIdsNumElts; ii += stride) - { - copyStruct.dstOutputIdsCBA[ii] = copyStruct.srcOutputIdsCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.logProbsNumElts; ii += stride) - { - copyStruct.dstLogProbsCBA[ii] = copyStruct.srcLogProbsCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.cumLogProbsNumElts; ii += stride) - { - copyStruct.dstCumLogProbs[ii] = copyStruct.srcCumLogProbs[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.sequenceLengthsNumElts; ii += stride) - { - copyStruct.dstSequenceLengthsCBA[ii] = copyStruct.srcSequenceLengthsCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.cumLogProbsCBANumElts; ii += stride) - { - copyStruct.dstCumLogProbsCBA[ii] = copyStruct.srcCumLogProbsCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.normedScoresNumElts; ii += stride) - { - copyStruct.dstNormedScoresCBA[ii] = copyStruct.srcNormedScoresCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.numBeamsNumElts; ii += stride) - { - copyStruct.dstNumBeamsCBA[ii] = copyStruct.srcNumBeamsCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.minNormedScoresNumElts; ii += stride) - { - copyStruct.dstMinNormedScoresCBA[ii] = copyStruct.srcMinNormedScoresCBA[ii]; - } - - for (SizeType32 ii = idx; ii < copyStruct.batchDonesNumElts; ii += stride) - { - copyStruct.dstBatchDones[ii] = copyStruct.srcBatchDones[ii]; - } -} - -void invokeCopyBeamHypotheses(DecodingOutput::BeamHypotheses const& src, DecodingOutput::BeamHypotheses const& dst, - ITensor& srcCumLogProbs, ITensor& dstCumLogProbs, runtime::CudaStream const& stream, SizeType32 numSMs) -{ - CopyBeamHypothesesStruct copyStruct = {}; - - copyStruct.srcOutputIdsCBA = bufferCast(*(src.outputIdsCBA)); - copyStruct.dstOutputIdsCBA = bufferCast(*(dst.outputIdsCBA)); - copyStruct.outputIdsNumElts = dst.outputIdsCBA->getSize(); - - copyStruct.srcLogProbsCBA = bufferCast(*(src.logProbsCBA)); - copyStruct.dstLogProbsCBA = bufferCast(*(dst.logProbsCBA)); - copyStruct.logProbsNumElts = dst.logProbsCBA->getSize(); - - copyStruct.srcSequenceLengthsCBA = bufferCast(*(src.sequenceLengthsCBA)); - copyStruct.dstSequenceLengthsCBA = bufferCast(*(dst.sequenceLengthsCBA)); - copyStruct.sequenceLengthsNumElts = dst.sequenceLengthsCBA->getSize(); - - copyStruct.srcCumLogProbsCBA = bufferCast(*(src.cumLogProbsCBA)); - copyStruct.dstCumLogProbsCBA = bufferCast(*(dst.cumLogProbsCBA)); - copyStruct.cumLogProbsCBANumElts = dst.cumLogProbsCBA->getSize(); - - copyStruct.srcNormedScoresCBA = bufferCast(*(src.normedScoresCBA)); - copyStruct.dstNormedScoresCBA = bufferCast(*(dst.normedScoresCBA)); - copyStruct.normedScoresNumElts = dst.normedScoresCBA->getSize(); - - copyStruct.srcNumBeamsCBA = bufferCast(*(src.numBeamsCBA)); - copyStruct.dstNumBeamsCBA = bufferCast(*(dst.numBeamsCBA)); - copyStruct.numBeamsNumElts = dst.numBeamsCBA->getSize(); - - copyStruct.srcMinNormedScoresCBA = bufferCast(*(src.minNormedScoresCBA)); - copyStruct.dstMinNormedScoresCBA = bufferCast(*(dst.minNormedScoresCBA)); - copyStruct.minNormedScoresNumElts = dst.minNormedScoresCBA->getSize(); - - copyStruct.srcBatchDones = bufferCast(*(src.batchDones)); - copyStruct.dstBatchDones = bufferCast(*(dst.batchDones)); - copyStruct.batchDonesNumElts = dst.batchDones->getSize(); - - copyStruct.srcCumLogProbs = bufferCast(srcCumLogProbs); - copyStruct.dstCumLogProbs = bufferCast(dstCumLogProbs); - copyStruct.cumLogProbsNumElts = srcCumLogProbs.getSize(); - - copyBeamHypotheses<<>>(copyStruct); -} - -__global__ void initializeOutput( - TokenIdType* finalOutputIds, TokenIdType const* endIds, SizeType32 const beam, SizeType32 const nMaxSeqLen) -{ - for (int i = threadIdx.x; i < nMaxSeqLen; i += blockDim.x) - { - finalOutputIds[blockIdx.x * nMaxSeqLen + i] = endIds[blockIdx.x / beam]; - } -} - -void invokeInitializeOutput(TokenIdType* finalOutputIds, TokenIdType const* endIds, SizeType32 const batch, - SizeType32 const beam, SizeType32 const nMaxSeqLen, cudaStream_t stream) -{ - initializeOutput<<>>(finalOutputIds, endIds, beam, nMaxSeqLen); -} - -__global__ void copyNextStepIds(TokenIdType* nextStepIds, TokenIdType const* const* outputIdsPtr, - SizeType32 const* sequenceLengths, SizeType32 const* numNewTokens, SizeType32 const* batchSlots, - SizeType32 batchSize, SizeType32 maxBatchSize, SizeType32 beamWidth, SizeType32 maxSeqLen, - SizeType32 maxTokensPerStep) -{ - for (auto index = static_cast(blockIdx.x * blockDim.x + threadIdx.x); - index < batchSize * beamWidth * maxTokensPerStep; index += static_cast(blockDim.x * gridDim.x)) - { - // numNewTokens == nullptr when Medusa is disabled - auto const batchIdx{index / (beamWidth * maxTokensPerStep)}; - auto const batchSlot{batchSlots[batchIdx]}; - auto const remainder{index % (beamWidth * maxTokensPerStep)}; - auto const beamIdx{remainder / maxTokensPerStep}; - auto const tokenIdx{remainder % maxTokensPerStep}; - auto const newTokens{numNewTokens == nullptr ? 1 : numNewTokens[batchSlot]}; - auto const batchBeamIdx = batchSlot * beamWidth + beamIdx; - auto const tokenBatchBeamIdx = tokenIdx * maxBatchSize * beamWidth + batchSlot * beamWidth + beamIdx; - auto const indexSrc = sequenceLengths[batchBeamIdx] - newTokens + tokenIdx; - if (tokenIdx >= newTokens || indexSrc < 0) - { - continue; - } - nextStepIds[tokenBatchBeamIdx] = outputIdsPtr[batchSlot][beamIdx * maxSeqLen + indexSrc]; - } -} - -void invokeCopyNextStepIds(TokenIdType* nextStepIds, TokenIdType const* const* outputIdsPtr, - SizeType32 const* sequenceLengths, SizeType32 const* numNewTokens, SizeType32 const* batchSlots, - SizeType32 batchSize, SizeType32 maxBatchSize, SizeType32 beamWidth, SizeType32 maxSeqLen, - SizeType32 maxTokensPerStep, cudaStream_t stream) -{ - int const numElems = batchSize * beamWidth * maxTokensPerStep; - dim3 block(min(256, numElems)); - dim3 grid(divUp(numElems, block.x)); - copyNextStepIds<<>>(nextStepIds, outputIdsPtr, sequenceLengths, numNewTokens, batchSlots, - batchSize, maxBatchSize, beamWidth, maxSeqLen, maxTokensPerStep); -} - -__global__ void transposeLogProbs(float* outputLogProbs, float* outputLogProbsTiled, SizeType32 const* sequenceLengths, - SizeType32 const* batchSlots, SizeType32 batchSize, SizeType32 maxBatchSize, SizeType32 beamWidth, - SizeType32 maxSeqLen) -{ - auto index = static_cast(blockIdx.x * blockDim.x + threadIdx.x); - - auto const batchIdx = index / (beamWidth * maxSeqLen); - auto const tmpIdx = index % (beamWidth * maxSeqLen); - auto const beamIdx = tmpIdx / maxSeqLen; - auto const pos = tmpIdx % maxSeqLen; - if (batchIdx >= batchSize) - { - return; - } - - auto const batchSlot = batchSlots[batchIdx]; - auto const batchBeamIdx = batchSlot * beamWidth + beamIdx; - if (pos < sequenceLengths[batchBeamIdx]) - { - auto const outputIndex = batchSlot * beamWidth * maxSeqLen + beamIdx * maxSeqLen + pos; - outputLogProbs[outputIndex] - = outputLogProbsTiled[pos * maxBatchSize * beamWidth + batchSlot * beamWidth + beamIdx]; - } -} - -void invokeTransposeLogProbs(float* outputLogProbs, float* outputLogProbsTiled, SizeType32 const* sequenceLengths, - SizeType32 const* batchSlots, SizeType32 batchSize, SizeType32 maxBatchSize, SizeType32 beamWidth, - SizeType32 maxSeqLen, cudaStream_t stream) -{ - dim3 block(256); - dim3 grid(divUp(batchSize * beamWidth * maxSeqLen, block.x)); - transposeLogProbs<<>>(outputLogProbs, outputLogProbsTiled, sequenceLengths, batchSlots, - batchSize, maxBatchSize, beamWidth, maxSeqLen); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END - -namespace tensorrt_llm::runtime::kernels -{ -// Must be similar to [cpp/tensorrt_llm/thop/gatherTreeOp.cpp] gatherTree -void gatherTree(DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, - SamplingConfig const& samplingConfig, runtime::CudaStream const& cudaStream, runtime::SizeType32 batchSlot) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& stream = cudaStream.get(); - BufferManager manager{std::make_shared(stream)}; - - auto& finalOutputIds = *decodingOutput.gatheredIds; - auto const& finalOutputIdsShape = finalOutputIds.getShape(); - auto const& decodingOutputIdsShape = decodingOutput.ids->getShape(); - auto const batchSize = finalOutputIdsShape.d[0]; - auto const beamWidth = finalOutputIdsShape.d[1]; - auto const maxSeqLength = finalOutputIdsShape.d[2]; - - TLLM_CHECK_WITH_INFO(beamWidth > 1, "gatherTree is only needed for beam search."); - - TLLM_CHECK_WITH_INFO(decodingOutputIdsShape.d[0] == batchSize, - common::fmtstr("Decoder batch size (" FMT_DIM ") does not match final batch size (" FMT_DIM ")", - decodingOutputIdsShape.d[0], batchSize)); - TLLM_CHECK_WITH_INFO(decodingOutputIdsShape.d[1] == beamWidth, - common::fmtstr("Decoder beam width (" FMT_DIM ") does not match final beam width (" FMT_DIM ")", - decodingOutputIdsShape.d[1], beamWidth)); - TLLM_CHECK_WITH_INFO(decodingOutputIdsShape.d[2] <= maxSeqLength, - common::fmtstr("Decoder seq length size (" FMT_DIM ") is too large for final seq length (" FMT_DIM ")", - decodingOutputIdsShape.d[2], maxSeqLength)); - - // prefill finalOutputIds with the EOS tokens from decodingInput.endIds - tensorrt_llm::kernels::invokeInitializeOutput(bufferCast(finalOutputIds), - bufferCast(*decodingInput.endIds), batchSize, beamWidth, maxSeqLength, stream); - sync_check_cuda_error(stream); - - std::vector lengthPenaltyVec; - auto lengthPenaltyPtr = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize}), TRTDataType::value)); - if (!samplingConfig.lengthPenalty.has_value() || samplingConfig.lengthPenalty.value().size() == 0) - { - lengthPenaltyVec = std::vector(batchSize, 1.0f); - } - else if (long int const size = samplingConfig.lengthPenalty.value().size(); size == 1) - { - lengthPenaltyVec = std::vector(batchSize, samplingConfig.lengthPenalty.value()[0]); - } - else - { - TLLM_CHECK_WITH_INFO(size == batchSize, - common::fmtstr("Size of lengthPenalty in SamplingConfig (" FMT_DIM ") is different from batchSize (" FMT_DIM - ")", - size, batchSize)); - lengthPenaltyVec = samplingConfig.lengthPenalty.value(); - } - - lengthPenaltyPtr = manager.copyFrom(lengthPenaltyVec, ITensor::makeShape({batchSize}), runtime::MemoryType::kGPU); - - tensorrt_llm::kernels::BeamHypotheses bh; - // logProbsTiled has shape [MSL, maxNumSequences, BM] and is passed unsliced. - // nMaxBatchSize must equal the allocation stride (dim-1), not the per-slot batchSize=1. - // The pointer is pre-offset by batchSlot*BM so that insertUnfinishedPathKernel, - // which uses bid=0 / nBatchSize=1, computes: - // (base + batchSlot*BM)[step * maxBS * BM + 0*BM + beamIdx] - // = base[step * maxBS * BM + batchSlot * BM + beamIdx] - // = logProbsTiled[step][batchSlot][beamIdx] ✓ - auto const logProbsTiledMaxBatchSize = static_cast(decodingOutput.logProbsTiled->getShape().d[1]); - auto const logProbsTiledBeamWidth = static_cast(decodingOutput.logProbsTiled->getShape().d[2]); - TLLM_CHECK_WITH_INFO(batchSlot < logProbsTiledMaxBatchSize, - "batchSlot (%d) must be < logProbsTiled maxBatchSize (%d); " - "logProbsTiled would be accessed out of bounds.", - batchSlot, logProbsTiledMaxBatchSize); - TLLM_CHECK_WITH_INFO(beamWidth == logProbsTiledBeamWidth, - "beamWidth (%d) must equal logProbsTiled BM dimension (%d); " - "pointer offset batchSlot*beamWidth would be misaligned.", - beamWidth, logProbsTiledBeamWidth); - bh.nMaxBatchSize = logProbsTiledMaxBatchSize; - bh.nBatchSize = batchSize; - bh.nBeamWidth = beamWidth; - bh.nMaxSeqLen = maxSeqLength; - bh.lengthPenalties = bufferCast(*lengthPenaltyPtr); - bh.inputLengths = bufferCast(*decodingInput.lengths); - bh.outputIds = bufferCast(finalOutputIds); - bh.logProbs = bufferCastOrNull(decodingOutput.logProbs); - bh.logProbsTiled = bufferCast(*decodingOutput.logProbsTiled) + batchSlot * beamWidth; - bh.sequenceLengths = bufferCast(*decodingOutput.lengths); - bh.cumLogProbs = bufferCast(*decodingOutput.cumLogProbs); - bh.outputIdsCBA = bufferCast(*decodingOutput.beamHypotheses.outputIdsCBA); - bh.logProbsCBA = bufferCast(*decodingOutput.beamHypotheses.logProbsCBA); - bh.sequenceLengthsCBA = bufferCast(*decodingOutput.beamHypotheses.sequenceLengthsCBA); - bh.cumLogProbsCBA = bufferCast(*decodingOutput.beamHypotheses.cumLogProbsCBA); - bh.normedScoresCBA = bufferCast(*decodingOutput.beamHypotheses.normedScoresCBA); - bh.numBeamsCBA = bufferCast(*decodingOutput.beamHypotheses.numBeamsCBA); - bh.minNormedScoresCBA = bufferCast(*decodingOutput.beamHypotheses.minNormedScoresCBA); - bh.batchDones = bufferCast(*decodingOutput.beamHypotheses.batchDones); - bh.finished = bufferCast(*decodingOutput.finishReasons); - bh.outputIdsUnfinish = bufferCast(*decodingOutput.ids); - bh.parentIdsUnfinish = bufferCast(*decodingOutput.parentIds); - - // This is where transpose is done - tensorrt_llm::kernels::invokeInsertUnfinishedPath(bh, stream); - sync_check_cuda_error(stream); - - tensorrt_llm::kernels::invokeFinalize(bh, stream); - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::runtime::kernels diff --git a/cpp/tensorrt_llm/kernels/decodingKernels.h b/cpp/tensorrt_llm/kernels/decodingKernels.h deleted file mode 100644 index 25fca71ee267..000000000000 --- a/cpp/tensorrt_llm/kernels/decodingKernels.h +++ /dev/null @@ -1,137 +0,0 @@ -/* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decodingInput.h" -#include "tensorrt_llm/runtime/decodingOutput.h" -#include "tensorrt_llm/runtime/samplingConfig.h" -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -struct gatherTreeParam -{ - // TODO rename the parameters - int32_t* beams = nullptr; // [batchSize, beamWidth, maxSeqLen], workspace to put intermediate outputIds - int32_t* sequenceLengths = nullptr; // [batchSize, beamWidth], total lengths of each query - int32_t maxSequenceLengthFinalStep = 0; - int32_t const* inputLengths = nullptr; // [batchSize, beamWidth] - // response input lengths (used to slice the ids during postprocessing) - int32_t* responseInputLengths = nullptr; - int32_t maxSeqLen = 0; - int32_t batchSize = 0; - int32_t beamWidth = 0; - int32_t const* stepIds = nullptr; // [maxSeqLen, batchSize, beamWidth] - int32_t const* parentIds = nullptr; // [maxSeqLen, batchSize, beamWidth] - int32_t const* endTokens = nullptr; // [batchSize], end token ids of each query - int32_t* outputIds = nullptr; // the buffer to put finalized ids - cudaStream_t stream; - float* cumLogProbs = nullptr; // [batchSize, beamWidth] - float lengthPenalty = 1.0f; - int earlyStopping = 1; -}; - -/* -Do gatherTree on beam search to get final result. -*/ -void invokeGatherTree(gatherTreeParam param); - -void invokeInsertUnfinishedPath(BeamHypotheses& bh, cudaStream_t stream); - -void invokeFinalize(BeamHypotheses& bh, cudaStream_t stream); - -//! \brief invoke the kernel that Initializes the output tensor by prefilling it with end tokens. -//! -//! \param finalOutputIds The output tensor to be initialized. -//! \param endIds The tensor containing the end IDs. -//! \param batchBeam batchSize*beamWidth. inferred from finalOutputIds.shape[0] * finalOutputIds.shape[1] -//! \param maxSeqLen The maximum sequence length, inferred from the finalOutputIds.shape[3] -//! \param stream The CUDA stream on which to perform the operation. -void invokeInitializeOutput(runtime::TokenIdType* finalOutputIds, runtime::TokenIdType const* endIds, - runtime::SizeType32 batch, runtime::SizeType32 beam, runtime::SizeType32 maxSeqLen, cudaStream_t stream); - -//! \brief Copies the data from the buffers in src to dst to reduce the kernel launch overhead of individual memcpy. -//! for streaming + beam search, where we need to avoid overwriting the beam search buffers. -//! -//! \param src the source, usually the buffers in which the beam search kernels write -//! \param dst temp buffers for use in the subsequent gatherTree kernels. -//! \param srcCumLogProbs source of the cumLogProbs. Separate since it's not included in beamHypotheses. -//! \param dstCumLogProbs dst of srcCumLogProbs. -//! \param stream CUDA stream to execute the kernel -//! \param numSMs number of SMs available on the device -void invokeCopyBeamHypotheses(runtime::DecodingOutput::BeamHypotheses const& src, - runtime::DecodingOutput::BeamHypotheses const& dst, runtime::ITensor& srcCumLogProbs, - runtime::ITensor& dstCumLogProbs, runtime::CudaStream const& stream, int numSMs); - -//! \brief Copies last numNewTokens (or 1 if numNewTokens == nullptr) tokens from outputIdsPtr -//! to nextStepIds according to sequenceLengths. -//! -//! \param nextStepIds output buffer [maxTokensPerStep, maxBatchSize, maxBeamWidth], -//! destination of the new tokens. -//! \param outputIdsPtr input buffer [maxBatchSize][maxBeamWidth, maxSeqLen], -//! array of pointers to the source of the copy. -//! \param sequenceLengths input buffer [maxBatchSize], sequence length of the request -//! in outputIdsPtr that includes all new tokens. It must be guaranteed that sequenceLengths <= maxSeqLen. -//! \param numNewTokens input buffer [maxBatchSize], optional, number of tokens to be copied. -//! If nullptr, only 1 token is copied. It must be guaranteed that numNewTokens <= sequenceLengths. -//! \param batchSlots input buffer [batchSize], address map from local index -//! to global index [0, batchSize] -> [0, maxBatchSize] -//! \param batchSize current batch size -//! \param maxBatchSize maximum batch size -//! \param beamWidth current beam width -//! \param maxSeqLen maximum sequence length -//! \param maxTokensPerStep maximum tokens per step -//! \param stream stream -void invokeCopyNextStepIds(runtime::TokenIdType* nextStepIds, runtime::TokenIdType const* const* outputIdsPtr, - runtime::SizeType32 const* sequenceLengths, runtime::SizeType32 const* numNewTokens, - runtime::SizeType32 const* batchSlots, runtime::SizeType32 batchSize, runtime::SizeType32 maxBatchSize, - runtime::SizeType32 beamWidth, runtime::SizeType32 maxSeqLen, runtime::SizeType32 maxTokensPerStep, - cudaStream_t stream); - -void invokeTransposeLogProbs(float* output_log_probs, float* output_log_probs_tiled, - runtime::SizeType32 const* sequence_lengths, runtime::SizeType32 const* batchSlots, runtime::SizeType32 batch_size, - runtime::SizeType32 max_batch_size, runtime::SizeType32 beam_width, runtime::SizeType32 max_seq_len, - cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END - -namespace tensorrt_llm::runtime::kernels -{ -//! \brief Inserts the running beams into the finished beams stored in the CBA buffers. (beams where the most likely -//! continuation is the end token get stored separately, and another candidate next token is stored). Then sorts the -//! beams according to their cumulative log probs. Note: the kernels in gatherTree modify the buffers inplace. When -//! streaming, we use tmp buffers since beam search kernels expect ungathered data. -//! -//! \param decodingOutput contains a slice of the output buffers to gather. Also contains the -//! DecodingOutput::BeamHypotheses object with the finished beams. -//! \param decodingInput used for endIds and input lengths. -//! \param samplingConfig the usual buffer samplingConfig. -//! \param cudaStream the CUDA stream on which to perform the operation. - -void gatherTree(DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, - SamplingConfig const& samplingConfig, runtime::CudaStream const& cudaStream, runtime::SizeType32 batchSlot = 0); -} // namespace tensorrt_llm::runtime::kernels diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.cu deleted file mode 100644 index 2f5eeb2c0a3c..000000000000 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.cu +++ /dev/null @@ -1,323 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaTypeUtils.cuh" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" - -#include "tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h" -#ifndef CUDART_VERSION -#error CUDART_VERSION Undefined! -#elif (CUDART_VERSION >= 11050) -#include -#else -#include "3rdparty/cub/cub.cuh" -#endif - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::runtime; - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::speculative_decoding -{ -namespace -{ - -template -__global__ void maskTargetLogitsKernel(T* targetLogits, SizeType32 const* batchSlots, SizeType32 beamWidth, - SizeType32 vocabSize, FinishedState const* finishedInput, SizeType32 maxBatchSize, - SizeType32* outputIdsAfterSampling, SizeType32* runtimeTopKDevicePtr, bool* maskBuffer) -{ - /** - * @brief Masking the selected token to -inf as was done in Huggingface TopK/TopP Logits Warper - * https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/generation/logits_process.py#L533 - */ - - auto const bid = blockIdx.x; - auto const batchIdx = bid / beamWidth; - auto const tid = static_cast(threadIdx.x); - auto const batchSlot = batchSlots[batchIdx]; - - constexpr bool IS_HALF = std::is_same::value; - T const MAX_T_VAL = (IS_HALF) ? HALF_FLT_MAX : FLT_MAX; - - auto targetLogitsBatch = targetLogits + batchIdx * vocabSize; - auto& finishedState = finishedInput[batchSlot]; - - auto* outputIdsAfterSamplingPtr = outputIdsAfterSampling + batchSlot * vocabSize; - auto* maskBufferBatch = maskBuffer + batchSlot * vocabSize; - - if (finishedState.isSkipDecoding() || finishedState.isFinished()) - { - return; - } - - __shared__ SizeType32 tokensToMask; - - if (tid == 0) - { - tokensToMask = runtimeTopKDevicePtr[batchSlot]; - } - __syncthreads(); - - for (SizeType32 vIdx = tid; vIdx < vocabSize; vIdx += static_cast(blockDim.x)) - { - if (outputIdsAfterSamplingPtr[vIdx] == -1) - { // we need to find the -1 boundary from returnAllTopP outputIds if topK == 0 or number of topP indices < topK - tokensToMask = vIdx; - } - maskBufferBatch[vIdx] = false; - } - - __syncthreads(); - if (tid == 0 && tokensToMask == 0) - { - // all tokens are selected if topK == 0 && topP ~= 1.0f - // in this case tokensToMask = vocabSize - tokensToMask = vocabSize; - } - __syncthreads(); - - for (SizeType32 vIdx = tid; vIdx < tokensToMask; vIdx += static_cast(blockDim.x)) - { - auto tokenToMask = outputIdsAfterSamplingPtr[vIdx]; - maskBufferBatch[tokenToMask] = true; - } - - __syncthreads(); - - for (SizeType32 vIdx = tid; vIdx < vocabSize; vIdx += static_cast(blockDim.x)) - { - if (!maskBufferBatch[vIdx]) - { - targetLogitsBatch[vIdx] = -MAX_T_VAL; - } - } -} - -template -__global__ void acceptDraftTokensKernel(T const* draftProbs, T* targetProbs, SizeType32 const* numsDraftTokens, - bool const* batchUseDraftLogits, TokenIdType const* draftIds, FinishedState const* finishedInput, - FinishedState* finishedOutput, curandState_t* curandState, SizeType32 const* batchSlots, SizeType32 maxDraftTokens, - SizeType32 beamWidth, SizeType32 vocabSize, bool randomThreshold, float constantThreshold, SizeType32 step, - bool* batchIsAccepted, SizeType32* targetOutputIds) -{ - auto const bid = blockIdx.x; - auto const draftTokenIdx = step; - auto const batchIdx = bid / beamWidth; - auto const beamIdx = bid % beamWidth; - auto const batchSlot = batchSlots[batchIdx]; - auto const batchSlotBeamWidth = batchSlot * beamWidth + beamIdx; - auto const tid = static_cast(threadIdx.x); - - auto const numDraftTokens = numsDraftTokens[batchSlotBeamWidth]; - auto const useDraftLogits = batchUseDraftLogits[batchSlotBeamWidth]; - - if (numDraftTokens == 0 || draftTokenIdx > numDraftTokens || finishedInput[batchSlot].isSkipDecoding() - || finishedInput[batchSlot].isFinished()) - { - if (tid == 0) - { - batchIsAccepted[batchSlot] = true; - - // either finished or skip decode in previous step, this step don't need decoding - finishedOutput[batchSlot].setSkipDecoding(); - - // if previous step is finished, write the state to next step too - if (finishedInput[batchSlot].isFinished()) - { - finishedOutput[batchSlot] = finishedInput[batchSlot]; - } - } - return; - } - - if (draftTokenIdx == numDraftTokens) - { - if (tid == 0) - { - batchIsAccepted[batchSlot] = false; - finishedOutput[batchSlot].setSkipDecoding(); - } - return; - } - // else (draftTokenIdx < numDraftTokens) - - auto const logitsOffset = (batchSlot * maxDraftTokens + draftTokenIdx) * beamWidth * vocabSize; - auto const draftProbsBatch = draftProbs + logitsOffset; - auto const targetProbsBatch = targetProbs + (batchIdx * beamWidth * vocabSize); - - __shared__ bool isAccepted; - __shared__ T sSumVal; - if (tid == 0) - { - auto const draftOutputTokenId = draftIds[batchSlot * maxDraftTokens + draftTokenIdx]; - if (useDraftLogits) - { - float threshold = randomThreshold ? curand_uniform(curandState + batchSlot) : constantThreshold; - auto const targetProb = static_cast(targetProbsBatch[draftOutputTokenId]); - auto const draftProb = static_cast(draftProbsBatch[draftOutputTokenId]); - isAccepted = targetProb > threshold * draftProb; - } - else - { - // Check if draft tokens are the same as target tokens - isAccepted = targetOutputIds[batchSlot] == draftOutputTokenId; - } - if (!isAccepted) - { - finishedOutput[batchSlot].setSkipDecoding(); - } - batchIsAccepted[batchSlot] = isAccepted; - } - - __syncthreads(); - - if (useDraftLogits && !isAccepted) - { - // correct target distribution - T const zeroVal = static_cast(0.0F); - T sumVal = zeroVal; - for (SizeType32 vIdx = tid; vIdx < vocabSize; vIdx += static_cast(blockDim.x)) - { - targetProbsBatch[vIdx] -= draftProbsBatch[vIdx]; - targetProbsBatch[vIdx] = targetProbsBatch[vIdx] >= zeroVal ? targetProbsBatch[vIdx] : zeroVal; - sumVal += targetProbsBatch[vIdx]; - } - sumVal = blockReduceSum(sumVal); - if (tid == 0) - { - sSumVal = sumVal; - } - __syncthreads(); - - for (SizeType32 vIdx = tid; vIdx < vocabSize; vIdx += static_cast(blockDim.x)) - { - targetProbsBatch[vIdx] /= sSumVal; - } - } -} - -__global__ void forwardAcceptedTokensKernel(SizeType32 batchSize, SizeType32 const* batchSlots, bool* batchIsAccepted, - SizeType32* sequenceLengths, TokenIdType const* draftIds, TokenIdType** idsPtrs, SizeType32 step, - SizeType32 maxDraftTokens, TokenIdType const* endIds, FinishedState* finishedOutput) -{ - auto index = static_cast(blockIdx.x * blockDim.x + threadIdx.x); - for (SizeType32 bi = index; bi < batchSize; bi += static_cast(gridDim.x * blockDim.x)) - { - auto const batchSlot = batchSlots[bi]; - if (batchIsAccepted[batchSlot] && !finishedOutput[batchSlot].isSkipDecoding() - && !finishedOutput[batchSlot].isFinished()) - { - auto const curSeqLen = sequenceLengths[batchSlot]; - auto const draftTokenIdx = step; - auto const draftOutputTokenId = draftIds[batchSlot * maxDraftTokens + draftTokenIdx]; - auto* outputIdsRequestPtr = idsPtrs[batchSlot]; - auto const outIdx = curSeqLen; - outputIdsRequestPtr[outIdx] = draftOutputTokenId; - if (outputIdsRequestPtr[outIdx] == endIds[batchSlot]) - { - finishedOutput[batchSlot].setFinishedEOS(); - // Do not increase seq len when EOS is generated. Seq len should always contain only tokens to be - // outputted - } - else - { - // We don't need to set output finished state as it is assumed to be in non finished state - sequenceLengths[batchSlot] += 1; - } - } - } -} // namespace - -} // namespace - -template -void invokeMaskTargetLogits(SizeType32 batchSize, T* targetLogits, SizeType32 const* batchSlots, SizeType32 beamWidth, - SizeType32 vocabSizePadded, FinishedState const* finishedInput, SizeType32 maxBatchSize, - SizeType32* outputIdsAfterSampling, SizeType32* runtimeTopKDevicePtr, bool* maskBuffer, cudaStream_t stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(beamWidth == 1); - { - dim3 block(1024); - dim3 grid(batchSize * beamWidth); - maskTargetLogitsKernel<<>>(targetLogits, batchSlots, beamWidth, vocabSizePadded, - finishedInput, maxBatchSize, outputIdsAfterSampling, runtimeTopKDevicePtr, maskBuffer); - } - sync_check_cuda_error(stream); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void invokeAcceptDraftTokens(SizeType32 batchSize, T* draftProbs, T* targetProbs, SizeType32 const* numsDraftTokens, - bool const* batchUseDraftLogits, TokenIdType const* draftIds, FinishedState const* finishedInput, - FinishedState* finishedOutput, curandState_t* curandState, SizeType32 const* batchSlots, SizeType32 maxDraftTokens, - SizeType32 beamWidth, SizeType32 vocabSizePadded, bool randomThreshold, float constantThreshold, SizeType32 step, - bool* batchIsAccepted, SizeType32* targetOutputIds, cudaStream_t stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(beamWidth == 1); - { - dim3 block(1024); - dim3 grid(batchSize * beamWidth); - acceptDraftTokensKernel<<>>(draftProbs, targetProbs, numsDraftTokens, - batchUseDraftLogits, draftIds, finishedInput, finishedOutput, curandState, batchSlots, maxDraftTokens, - beamWidth, vocabSizePadded, randomThreshold, constantThreshold, step, batchIsAccepted, targetOutputIds); - } - sync_check_cuda_error(stream); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template void invokeMaskTargetLogits(SizeType32 batchSize, float* targetLogits, SizeType32 const* batchSlots, - SizeType32 beamWidth, SizeType32 vocabSizePadded, FinishedState const* finishedInput, SizeType32 maxBatchSize, - SizeType32* outputIdsAfterSampling, SizeType32* runtimeTopKDevicePtr, bool* maskBuffer, cudaStream_t stream); -template void invokeMaskTargetLogits(SizeType32 batchSize, half* targetLogits, SizeType32 const* batchSlots, - SizeType32 beamWidth, SizeType32 vocabSizePadded, FinishedState const* finishedInput, SizeType32 maxBatchSize, - SizeType32* outputIdsAfterSampling, SizeType32* runtimeTopKDevicePtr, bool* maskBuffer, cudaStream_t stream); - -template void invokeAcceptDraftTokens(SizeType32 batchSize, float* draftProbs, float* targetProbs, - SizeType32 const* numsDraftTokens, bool const* batchUseDraftLogits, TokenIdType const* draftIds, - FinishedState const* finishedInput, FinishedState* finishedOutput, curandState_t* curandState, - SizeType32 const* batchSlots, SizeType32 maxDraftTokens, SizeType32 beamWidth, SizeType32 vocabSizePadded, - bool randomThreshold, float constantThreshold, SizeType32 step, bool* batchIsAccepted, SizeType32* targetOutputIds, - cudaStream_t stream); -template void invokeAcceptDraftTokens(SizeType32 batchSize, half* draftProbs, half* targetProbs, - SizeType32 const* numsDraftTokens, bool const* batchUseDraftLogits, TokenIdType const* draftIds, - FinishedState const* finishedInput, FinishedState* finishedOutput, curandState_t* curandState, - SizeType32 const* batchSlots, SizeType32 maxDraftTokens, SizeType32 beamWidth, SizeType32 vocabSizePadded, - bool randomThreshold, float constantThreshold, SizeType32 step, bool* batchIsAccepted, SizeType32* targetOutputIds, - cudaStream_t stream); - -void invokeForwardAcceptedTokens(SizeType32 batchSize, SizeType32 const* batchSlots, bool* batchIsAccepted, - SizeType32* outputSequenceLengths, TokenIdType const* draftIds, TokenIdType** idsPtrs, SizeType32 step, - SizeType32 maxDraftTokens, TokenIdType const* endIds, FinishedState* finishedOutput, cudaStream_t stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - dim3 block(std::min(static_cast(batchSize), 256u)); - dim3 grid(divUp(static_cast(batchSize), block.x)); - forwardAcceptedTokensKernel<<>>(batchSize, batchSlots, batchIsAccepted, - outputSequenceLengths, draftIds, idsPtrs, step, maxDraftTokens, endIds, finishedOutput); - sync_check_cuda_error(stream); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} -} // namespace kernels::speculative_decoding - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h deleted file mode 100644 index 92fb3f68981f..000000000000 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h +++ /dev/null @@ -1,103 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/speculativeDecoding/common.h" -#include "tensorrt_llm/runtime/common.h" -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::speculative_decoding -{ - -//! \brief Accepts or rejects draft tokens based on their probability distributions or the equality of draft and target -//! tokens. Corrects targetLogits for the last accepted token -//! according to https://openreview.net/pdf?id=C9NEblP8vS -//! -//! \param batchSize current batch size -//! \param draftProbs output buffer [maxDraftTokens, batchSize, beamWidth, vocabSize]. -//! Workspace buffer for token probabilities of the draft model. -//! \param targetProbs output buffer [maxDraftTokens+1, batchSize, beamWidth, vocabSize]. -//! Workspace buffer for token probabilities of the target model. -//! \param numsDraftTokens input buffer [batchSize]. Number of draft tokens per request -//! \param batchUseDraftLogits input buffer [batchSize]. Acceptance logic using draft logits or not, per request -//! \param draftIds input buffer [batchSize, draftTokens]. Pointer to draft token ids. -//! \param finishedInput input buffer [batchSize, beamWidth]. -//! \param finishedOutput output buffer [batchSize, beamWidth]. At each step sets SKIP_DECODING if token is not -//! accepted. -//! \param curandState input buffer [batchSize]. Curand states properly initialized using invokeCurandInitialize -//! per request. -//! \param batchSlots input buffer [batchSize], address map from local index to global index [0, batchSize] -> -//! [0, maxBatchSize]. -//! \param maxDraftTokens maximum number of draft tokens -//! \param beamWidth beam width (only beamWidth == 1 supported) -//! \param vocabSizePadded padded vocab size -//! \param randomThreshold True if use uniformly sampled threshold for token acceptance -//! \param constantThreshold threshold used to accept tokens if randomThreshold is false -//! \param step The current step of decoding (draft token id index) -//! \param batchIsAccepted output buffer [batchSize]. Stores acceptance result for multinomial sampling later or -//! forwarding next step. -//! \param targetOutputIds input/output buffer [batchSize]. Stores target sampling output ids for acceptById -//! logics. -//! \param stream stream -template -void invokeAcceptDraftTokens(runtime::SizeType32 batchSize, T* draftProbs, T* targetProbs, - runtime::SizeType32 const* numsDraftTokens, bool const* batchUseDraftLogits, runtime::TokenIdType const* draftIds, - FinishedState const* finishedInput, FinishedState* finishedOutput, curandState_t* curandState, - runtime::SizeType32 const* batchSlots, runtime::SizeType32 maxDraftTokens, runtime::SizeType32 beamWidth, - runtime::SizeType32 vocabSizePadded, bool randomThreshold, float constantThreshold, runtime::SizeType32 step, - bool* batchIsAccepted, runtime::SizeType32* targetOutputIds, cudaStream_t stream); - -//! \brief Mask the target logits with -inf for unselected topK/topP token ids. -//! according to -//! https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/generation/utils.py#L4064 -//! -//! \param batchSize current batch size -//! \param targetLogits input/output buffer [batchSize][draftTokens+1, beamWidth, vocabSize]. -//! Vector of pointers to the logits. (beamWidth == 1) -//! Initially contains token logits of the target model. -//! \param batchSlots input buffer [batchSize], address map from local index to global index [0, batchSize] -> -//! [0, maxBatchSize]. -//! \param beamWidth beam width (only beamWidth == 1 supported) -//! \param vocabSizePadded padded vocab size -//! \param finishedInput input buffer [batchSize, beamWidth]. -//! \param maxBatchSize maximum batch size -//! \param outputIdsAfterSampling input buffer [batchSize, vocabSize]. Stores all selected IDs from sampling for -//! masking. -//! \param numsDraftTokens input buffer [batchSize]. Number of draft tokens per request -//! \param runtimeTopKDevicePtr input buffer [batchSize] the topks in sampling step, for porting topK ids out. -//! \param maskBuffer input buffer [batchSize, vocabSize] for masking calculation (index value to position). -//! \param stream stream -template -void invokeMaskTargetLogits(runtime::SizeType32 batchSize, T* targetLogits, runtime::SizeType32 const* batchSlots, - runtime::SizeType32 beamWidth, runtime::SizeType32 vocabSizePadded, FinishedState const* finishedInput, - runtime::SizeType32 maxBatchSize, runtime::SizeType32* outputIdsAfterSampling, - runtime::SizeType32* runtimeTopKDevicePtr, bool* maskBuffer, cudaStream_t stream); - -void invokeForwardAcceptedTokens(runtime::SizeType32 batchSize, runtime::SizeType32 const* batchSlots, - bool* batchIsAccepted, runtime::SizeType32* outputSequenceLengths, runtime::TokenIdType const* draftIds, - runtime::TokenIdType** idsPtrs, runtime::SizeType32 step, runtime::SizeType32 maxDraftTokens, - runtime::TokenIdType const* endIds, FinishedState* finishedOutput, cudaStream_t stream); - -} // namespace kernels::speculative_decoding - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h index fd3cbd53512f..3e1183fd77df 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h @@ -43,7 +43,7 @@ using IndexType = int; * @param rewindDraftTokenCount : Count to rewind * @param seqSlotRemapping mapping from batch index to index of the seqSlot in the sorted seqSlot buffer * e.g. for requests [0, 1, 2] with seqSlots [5, 3, 4], seqSlotRemapping is [1, 2, 0] - * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices from gptDecoderBatched + * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices produced by the decoder * and pointerArray and pastKeyValueLengths from runtimeBuffers. * @param maxKVCacheLen : Maximum length of each KV cache * @param stream : CUDA stream to use. @@ -69,7 +69,7 @@ void updateLinearKVCacheDraftTokenLocationCommonRewind(runtime::SizeType32 const * @param rewindDraftTokenCount : Count to rewind * @param seqSlotRemapping mapping from batch index to index of the seqSlot in the sorted seqSlot buffer * e.g. for requests [0, 1, 2] with seqSlots [5, 3, 4], seqSlotRemapping is [1, 2, 0] - * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices from gptDecoderBatched + * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices produced by the decoder * and pointerArray and pastKeyValueLengths from runtimeBuffers. * @param maxKVCacheLen : Maximum length of each KV cache * @param maxBlocksPerSeq : Maximum blocks per sequence of Block KV cache. @@ -99,7 +99,7 @@ void updateKVBlockArrayDraftTokenLocationCommonRewind(runtime::SizeType32 const* * one sequence. * @param seqSlotRemapping mapping from batch index to index of the seqSlot in the sorted seqSlot buffer * e.g. for requests [0, 1, 2] with seqSlots [5, 3, 4], seqSlotRemapping is [1, 2, 0] - * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices from gptDecoderBatched + * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices produced by the decoder * and pointerArray and pastKeyValueLengths from runtimeBuffers. * @param maxKVCacheLen : Maximum length of each KV cache * @param stream : CUDA stream to use. @@ -127,7 +127,7 @@ void updateLinearKVCacheDraftTokenLocationSeparateRewind(runtime::SizeType32 con * one sequence. * @param seqSlotRemapping mapping from batch index to index of the seqSlot in the sorted seqSlot buffer * e.g. for requests [0, 1, 2] with seqSlots [5, 3, 4], seqSlotRemapping is [1, 2, 0] - * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices from gptDecoderBatched + * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices produced by the decoder * and pointerArray and pastKeyValueLengths from runtimeBuffers. * @param maxKVCacheLen : Maximum length of each KV cache * @param maxBlocksPerSeq : Maximum blocks per sequence of Block KV cache. @@ -160,7 +160,7 @@ void updateKVBlockArrayDraftTokenLocationSeparateRewind(runtime::SizeType32 cons * rewind adjustment for one sequence. * @param seqSlotRemapping mapping from batch index to index of the seqSlot in the sorted seqSlot buffer * e.g. for requests [0, 1, 2] with seqSlots [5, 3, 4], seqSlotRemapping is [1, 2, 0] - * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices from gptDecoderBatched + * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices produced by the decoder * and pointerArray and pastKeyValueLengths from runtimeBuffers. * @param maxKVCacheLen : Maximum length of each KV cache * @param stream : CUDA stream to use. @@ -191,7 +191,7 @@ void updateLinearKVCacheDraftTokenLocation(runtime::SizeType32 const* seqAccepte * rewind adjustment for one sequence, indexed through batchSlots. * @param seqSlotRemapping mapping from batch index to index of the seqSlot in the sorted seqSlot buffer * e.g. for requests [0, 1, 2] with seqSlots [5, 3, 4], seqSlotRemapping is [1, 2, 0] - * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices from gptDecoderBatched + * Required to match seqAcceptedDraftTokenOffsets and packedAcceptedDraftTokensIndices produced by the decoder * and pointerArray and pastKeyValueLengths from runtimeBuffers. * @param batchSlots : [seqCount] indices of sequences in the seq slots. * @param maxKVCacheLen : Maximum length of each KV cache diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.cu deleted file mode 100644 index c109f28e9a32..000000000000 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.cu +++ /dev/null @@ -1,71 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaTypeUtils.cuh" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" - -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#ifndef CUDART_VERSION -#error CUDART_VERSION Undefined! -#elif (CUDART_VERSION >= 11050) -#include -#else -#include "3rdparty/cub/cub.cuh" -#endif - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::runtime; - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::speculative_decoding -{ -namespace -{ -__global__ void scatterMedusaDraftTokens(TokenIdType* treeDraftIds, TokenIdType const* sourceDraftIds, - SizeType32 const* treeIds, SizeType32 const* tokensPerStepData, SizeType32 const* batchSlots, - SizeType32 maxDecodingTokens) -{ - auto const batchIdx = static_cast(blockIdx.x); - auto const batchSlot = batchSlots[batchIdx]; - auto const tokensPerStep = tokensPerStepData[batchSlot]; - auto const maxDecodingDraftTokens = maxDecodingTokens - 1; - for (auto index = static_cast(threadIdx.x); index < tokensPerStep - 1; - index += static_cast(blockDim.x)) - { - auto const indexInTree = treeIds[batchSlot * maxDecodingDraftTokens + index]; - auto const treeDraftIdx = batchSlot * maxDecodingDraftTokens + index; - auto const sourceDraftIdx = batchSlot * maxDecodingTokens + indexInTree; - treeDraftIds[treeDraftIdx] = sourceDraftIds[sourceDraftIdx]; - } -} -} // namespace - -void scatterMedusaDraftTokens(TokenIdType* treeDraftIds, TokenIdType const* sourceDraftIds, SizeType32 const* treeIds, - SizeType32 const* tokensPerStep, SizeType32 const* batchSlots, SizeType32 maxDecodingTokens, SizeType32 batchSize, - cudaStream_t stream) -{ - constexpr SizeType32 BLOCK_SIZE = 256; - scatterMedusaDraftTokens<<>>( - treeDraftIds, sourceDraftIds, treeIds, tokensPerStep, batchSlots, maxDecodingTokens); -} -} // namespace kernels::speculative_decoding - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h deleted file mode 100644 index 8e79aa653ecb..000000000000 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h +++ /dev/null @@ -1,53 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/speculativeDecoding/common.h" -#include "tensorrt_llm/runtime/common.h" -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::speculative_decoding -{ - -//! \brief assembles draft tokens to treeDraftIds from sourceDraftIds using indices of treeIds -//! -//! \param treeDraftIds output buffer [maxBatchSize, maxDecodingTokens-1], output draft tokens -//! scattered from sourceDraftIds according to treeIds111 -//! \param sourceDraftIds input buffer [maxBatchSize, maxDecodingTokens], draft tokens saved leanearly after -//! sampling from Medusa heads with TopK. -//! \param treeIds input buffer [maxBatchSize, maxDecodingTokens-1], address map from sourceDraftIds to treeDraftIds -//! [0, unqiueDraftTokens] -> [0, maxDecodingTokens], where unqiueDraftTokens = sum(MedusaHeadsTopK) -//! unqiueDraftTokens <= maxDraftTokens -//! \param tokensPerStep input buffer [maxBatchSize], number of output draft tokens -//! \param batchSlots input buffer [maxBatchSize], address map from local index -//! to global index [0, batchSize] -> [0, maxBatchSize] -//! \param maxDecodingTokens maximum number of tokens per step configured in the system -//! \param batchSize current batch size -//! \param stream cuda stream -void scatterMedusaDraftTokens(runtime::TokenIdType* treeDraftIds, runtime::TokenIdType const* sourceDraftIds, - runtime::SizeType32 const* treeIds, runtime::SizeType32 const* tokensPerStep, runtime::SizeType32 const* batchSlots, - runtime::SizeType32 maxDecodingTokens, runtime::SizeType32 batchSize, cudaStream_t stream); - -} // namespace kernels::speculative_decoding - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/layers/CMakeLists.txt b/cpp/tensorrt_llm/layers/CMakeLists.txt deleted file mode 100644 index 9ce72c91da4e..000000000000 --- a/cpp/tensorrt_llm/layers/CMakeLists.txt +++ /dev/null @@ -1,37 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & -# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); you may not -# use this file except in compliance with the License. You may obtain a copy of -# the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT -# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the -# License for the specific language governing permissions and limitations under -# the License. -# -file(GLOB SRCS *.cpp) -file(GLOB CU_SRCS *.cu) - -if(NOT WIN32) - # additional warnings - # - # Ignore overloaded-virtual warning. We intentionally change parameters of - # some methods in derived class. - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall") - if(WARNING_IS_ERROR) - message(STATUS "Treating warnings as errors in GCC compilation") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror") - endif() -else() # Windows - # warning level 4 - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /W4") -endif() - -add_library(layers_src OBJECT ${SRCS} ${CU_SRCS}) -set_property(TARGET layers_src PROPERTY POSITION_INDEPENDENT_CODE ON) -set_property(TARGET layers_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) diff --git a/cpp/tensorrt_llm/layers/banWordsLayer.cpp b/cpp/tensorrt_llm/layers/banWordsLayer.cpp deleted file mode 100644 index 27bf22a28238..000000000000 --- a/cpp/tensorrt_llm/layers/banWordsLayer.cpp +++ /dev/null @@ -1,165 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "banWordsLayer.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/banBadWords.h" -#include "tensorrt_llm/kernels/banRepeatNgram.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" - -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -BanWordsLayer::BanWordsLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - allocateBuffer(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BanWordsLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (mDecodingMode.isUseNoRepeatNgramSize()) - { - mNoRepeatNgramSizeDevice - = mBufferManager->gpu(ITensor::makeShape({mDecoderDomain.getBatchSize()}), TRTDataType::value); - } - - mNoRepeatNgramSize = mBufferManager->pinnedPool( - ITensor::makeShape({mDecoderDomain.getBatchSize()}), TRTDataType::value); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BanWordsLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(BanWordsLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - auto const& banWordsParams = setupParams->banWordsParams; - TLLM_CHECK_WITH_INFO(banWordsParams, "banWordsParams for setup is not set"); - bool const useNoRepeatNgramSize - = mDecodingMode.isUseNoRepeatNgramSize() && banWordsParams->noRepeatNgramSize.has_value(); - FillBuffers const fillBuffers{batchSize, mDecoderDomain.getBatchSize(), mBufferManager}; - mUseNoRepeatNgramSize |= useNoRepeatNgramSize; - if (mUseNoRepeatNgramSize) - { - fillBuffers(banWordsParams->noRepeatNgramSize, DefaultDecodingParams::getNoRepeatNgramSize(), - mNoRepeatNgramSize, mNoRepeatNgramSizeDevice, batchSlots, - std::make_pair(0.f, std::numeric_limits::max()), "no_repeat_ngram_size"); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BanWordsLayer::banRepeatNGrams(TensorPtr const& logits, std::shared_ptr const& outputs, - std::shared_ptr const& inputs, BufferConstPtr const& batchSlots, BufferPtr noRepeatNgramSizeDevice, - DecoderDomain const& decoderDomain, SizeType32 maxSeqLen, bool useNoRepeatNgramSize) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - if (useNoRepeatNgramSize) - { - // auto const maxStep = inputs->step; // TODO Should we use step? but current inputs->step is always 0. - auto const maxStep = maxSeqLen; - // Temporary variables to store dereferenced inputs - auto logitsPtr = bufferCast(*logits); - auto outputIdsPtr = bufferCast(*outputs->outputIdsPtr); - auto finishedPtr - = reinterpret_cast(bufferCastOrNull(inputs->finished)); - auto parentIdsPtr = bufferCast(*outputs->parentIdsPtr); - auto batchSlotsPtr = bufferCast(*batchSlots); - auto sequenceLengthPtr = bufferCast(*outputs->sequenceLength.value()); - auto noRepeatNgramSizeDevicePtr = bufferCastOrNull(noRepeatNgramSizeDevice); - - // Call to invokeBanRepeatNgram with dereferenced inputs - invokeBanRepeatNgram(logitsPtr, outputIdsPtr, finishedPtr, parentIdsPtr, batchSlotsPtr, sequenceLengthPtr, - decoderDomain.getBatchSize(), decoderDomain.getBeamWidth(), maxSeqLen, noRepeatNgramSizeDevicePtr, - decoderDomain.getVocabSizePadded(), maxStep, getStream()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BanWordsLayer::banBadWords(TensorPtr const& logits, std::shared_ptr const& outputs, - std::shared_ptr const& inputs, BufferConstPtr const& batchSlots, DecoderDomain const& decoderDomain, - SizeType32 maxSeqLen) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const maxBadWordsLength = inputs->banWordsInputs->maxBadWordsLen; - if (maxBadWordsLength != 0) - { - // Temporary variables to store dereferenced inputs - auto badWordsPtr = bufferCast(*inputs->banWordsInputs->badWordsPtr.value()); - auto badWordsLens = bufferCast(*inputs->banWordsInputs->badWordsLengths.value()); - auto logitsPtr = bufferCast(*logits); - auto outputIdsPtr = bufferCast(*outputs->outputIdsPtr); - auto parentIdsPtr - = decoderDomain.getBeamWidth() > 1 ? bufferCast(*outputs->parentIdsPtr) : nullptr; - auto sequenceLengthPtr = bufferCast(*outputs->sequenceLength.value()); - auto batchSlotsPtr = bufferCast(*batchSlots); - - // Call to invokeBanBadWords with dereferenced inputs - invokeBanBadWords(logitsPtr, outputIdsPtr, parentIdsPtr, batchSlotsPtr, decoderDomain.getBatchSize(), - decoderDomain.getBeamWidth(), badWordsPtr, badWordsLens, maxBadWordsLength, - decoderDomain.getVocabSizePadded(), sequenceLengthPtr, maxSeqLen, getStream()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BanWordsLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(BanWordsLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - auto outputs = std::dynamic_pointer_cast(baseOutputs); - - TLLM_CHECK_WITH_INFO(inputs->banWordsInputs, "banWordsInputs for forward is not set"); - - auto const localDecoderDomain = getLocalDecoderDomain(inputs, mDecoderDomain); - auto const maxSeqLen = outputs->outputIds->getDimension<-1>(); - - banRepeatNGrams(workspace->getDeviceRuntimeLogits(), outputs, inputs, workspace->getDeviceBatchSlots(), - mNoRepeatNgramSizeDevice, localDecoderDomain, maxSeqLen, mUseNoRepeatNgramSize); - banBadWords(workspace->getDeviceRuntimeLogits(), outputs, inputs, workspace->getDeviceBatchSlots(), - localDecoderDomain, maxSeqLen); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class BanWordsLayer; -template class BanWordsLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/banWordsLayer.h b/cpp/tensorrt_llm/layers/banWordsLayer.h deleted file mode 100644 index 777c0e86f8ee..000000000000 --- a/cpp/tensorrt_llm/layers/banWordsLayer.h +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Layer to ban specific words from being sampled. -//! Supports banning bad words and repeating N grams. -//! Set badWordsPtr, maxBadWordsLen and badWordsLengths to ban bad words. -//! Set noRepeatNgramSize in input params to ban repeat Ngrams. -//! Layer modifies logits in-place. -template -class BanWordsLayer : public BaseLayer -{ - -public: - BanWordsLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) override; - - //! \brief Modifies 'outputs->logits' in-place with -INF for banned words - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - -private: - void allocateBuffer(); - void banBadWords(TensorPtr const& logits, std::shared_ptr const& outputs, - std::shared_ptr const& inputs, BufferConstPtr const& batchSlots, - DecoderDomain const& decoderDomain, runtime::SizeType32 maxSeqLen); - void banRepeatNGrams(TensorPtr const& logits, std::shared_ptr const& outputs, - std::shared_ptr const& inputs, BufferConstPtr const& batchSlots, - BufferPtr noRepeatNgramSizeDevice, DecoderDomain const& decoderDomain, runtime::SizeType32 maxSeqLen, - bool useNoRepeatNgramSize); - -private: - executor::DecodingMode mDecodingMode; - - TensorPtr mNoRepeatNgramSizeDevice; - TensorPtr mNoRepeatNgramSize; - bool mUseNoRepeatNgramSize{false}; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/baseLayer.h b/cpp/tensorrt_llm/layers/baseLayer.h deleted file mode 100644 index 5da28609f326..000000000000 --- a/cpp/tensorrt_llm/layers/baseLayer.h +++ /dev/null @@ -1,109 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include - -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" - -namespace tensorrt_llm::layers -{ - -class BaseLayer -{ -public: - using SizeType32 = runtime::SizeType32; - using TokenIdType = runtime::TokenIdType; - using BufferConstPtr = runtime::IBuffer::SharedConstPtr; - using BufferPtr = runtime::IBuffer::SharedPtr; - using TensorConstPtr = runtime::ITensor::SharedConstPtr; - using TensorPtr = runtime::ITensor::SharedPtr; - - BaseLayer(DecoderDomain decoderDomain, std::shared_ptr bufferManager) - : mBufferManager(std::move(bufferManager)) - , mDecoderDomain(std::move(decoderDomain)) - { - } - - virtual ~BaseLayer() = default; - - //! @returns cuda stream associated with layer - [[nodiscard]] cudaStream_t getStream() const noexcept - { - return mBufferManager->getStream().get(); - } - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] virtual size_t getWorkspaceSize() const noexcept - { - return 0; - }; - - // clang-format off - //! \brief Virtual function to setup internal states of the layer with sampling params - //! specified in setupParams for the entries specified by batchSlots. - //! It updates data for new requests in internal tensors inplace. - //! Thus, it must be called only once for new requests. - //! - //! \param batchSize current batch size configured in the system - //! \param beamWidth current beam width configured in the system - //! \param batchSlots input buffer [maxBatchSize], address map of the new requests, in pinned memory - //! \param setupParams shared pointer to params inherited from BaseSetupParams - // clang-format on - virtual void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) - = 0; - - // clang-format off - //! \brief Virtual function to execute layer async on GPU. - //! There must be no stream synchronization inside this function. - //! - //! \param outputs shared pointer to params inherited from BaseDecodingOutputs - //! \param inputs shared pointer to params inherited from BaseForwardParams - // clang-format on - virtual void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) - = 0; - - // clang-format off - //! \brief Virtual function to execute layer synchronously on CPU / GPU. - //! It is allowed (but not necassary) to synchronize on stream inside this function. - //! It is targeted mainly for prototyping. - //! - //! \param outputs shared pointer to params inherited from BaseDecodingOutputs - //! \param inputs shared pointer to params inherited from BaseForwardParams - // clang-format on - virtual void forwardSync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) - { - } - -protected: - // Buffer Manager - std::shared_ptr mBufferManager; - - // Domain in which token decoding is computed - DecoderDomain mDecoderDomain; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/beamSearchLayer.cu b/cpp/tensorrt_llm/layers/beamSearchLayer.cu deleted file mode 100644 index b288e9c0031b..000000000000 --- a/cpp/tensorrt_llm/layers/beamSearchLayer.cu +++ /dev/null @@ -1,424 +0,0 @@ -/* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "beamSearchLayer.h" -#include "tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h" - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" - -#include - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::kernels; - -namespace tensorrt_llm::layers -{ - -#define GET_INFO_STAGE1(paddedBeamWidth) \ - { \ - int constexpr nBlock = (paddedBeamWidth < 16) ? ((paddedBeamWidth < 8) ? kThreadForSmallBeamWidth : 128) : 64; \ - TLLM_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( \ - &nMaxActiveBlock, beamStage1Kernel, nBlock, 0)); \ - TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, beamStage1Kernel)); \ - break; \ - } - -#define GET_INFO_STAGE2(paddedBeamWidth) \ - { \ - if (nByteDynamicSharedMemoryStage2 > nByteMaxSharedMemoryPerBlock) \ - { \ - TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, beamStage2Kernel)); \ - } \ - else if (nVPart <= 32) \ - { \ - TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, beamStage2Kernel)); \ - } \ - else if (nVPart <= 64) \ - { \ - TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, beamStage2Kernel)); \ - } \ - else \ - { \ - TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, beamStage2Kernel)); \ - } \ - break; \ - } - -#define GET_INFO_STAGE3(paddedBeamWidth, isV2) \ - { \ - int constexpr nThreadStage3 = (paddedBeamWidth + 31) / 32 * 32; \ - TLLM_CUDA_CHECK( \ - cudaFuncGetAttributes(&attr, beamStage3Kernel)); \ - break; \ - } - -template -BeamSearchLayer::BeamSearchLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 const batchSize{mDecoderDomain.getBatchSize()}; - SizeType32 const beamWidth{mDecoderDomain.getBeamWidth()}; - SizeType32 const vocabSize{mDecoderDomain.getVocabSize()}; - TLLM_CHECK_WITH_INFO(beamWidth <= kMaxBeamWidth, "Beam width is larger than the maximum supported (%d > %d)", - int(beamWidth), int(kMaxBeamWidth)); - this->mVBWS = mode.isUseVariableBeamWidthSearch(); - - allocateBuffer(); - configureBeamSearchLayer(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BeamSearchLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 const batchSize{mDecoderDomain.getBatchSize()}; - auto const batchSizeShape{ITensor::makeShape({batchSize})}; - auto const batchSizeXBeamWidthArraySizeShape{ - ITensor::makeShape({batchSize * static_cast(kMaxBeamWidthArrayLength)})}; - - mBeamSearchDiversityRateHost = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mBeamSearchDiversityRateDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - mLengthPenaltyHost = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mLengthPenaltyDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - mEarlyStoppingHost = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mEarlyStoppingDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - if (this->mVBWS) - { - mBeamWidthArrayHost = mBufferManager->pinnedPool(batchSizeXBeamWidthArraySizeShape, TRTDataType::value); - mBeamWidthArrayDevice = mBufferManager->gpu(batchSizeXBeamWidthArraySizeShape, TRTDataType::value); - - mBeamWidthIn = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mBeamWidthOut = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BeamSearchLayer::configureBeamSearchLayer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 const batchSize{mDecoderDomain.getBatchSize()}; - SizeType32 const beamWidth{mDecoderDomain.getBeamWidth()}; - SizeType32 const vocabSize{mDecoderDomain.getVocabSize()}; - SizeType32 const paddedBeamWidth{padToNextPowerOfTwo(beamWidth)}; - cudaFuncAttributes attr; - - // Find device information to determine `nVPart`. - int const nByteMaxSharedMemoryPerSM = getMaxSharedMemoryPerSM(); - int const nByteMaxSharedMemoryPerBlock = getMaxSharedMemoryPerBlockOptin(); - int const nByteReservedSharedMemoryPerBlock = nByteMaxSharedMemoryPerSM - nByteMaxSharedMemoryPerBlock; - this->mByteMaxSharedMemoryPerBlock = nByteMaxSharedMemoryPerBlock; - - if (beamWidth <= kMaxBeamWidthForV1 && !(this->mVBWS)) - { - // V1 workflow for small beam width and non-VBWS - // Stage 1 - int nMaxActiveBlock = -1; - switch (paddedBeamWidth) - { - case 1: GET_INFO_STAGE1(1); - case 2: GET_INFO_STAGE1(2); - case 4: GET_INFO_STAGE1(4); - case 8: GET_INFO_STAGE1(8); - default: break; - } - int nByteStaticSharedMemory = attr.sharedSizeBytes; - int nByteMaxDynamicSharedMemoryPerBlock = nByteMaxSharedMemoryPerBlock - nByteStaticSharedMemory; - // Find the maximum of `nBlock` (maximum of `nVPart`, minimum of `nByteDynamicSharedMemoryStage1`), s.t. - // `nVPart <= kMaxVPartStage1 && nByteDynamicSharedMemoryStage1 * nVPart >= sizeof(T) * vocabSize` - TLLM_CHECK_WITH_INFO(nByteMaxDynamicSharedMemoryPerBlock * kMaxVPartStage1 >= sizeof(T) * vocabSize, - "vocab_size is too large for Beam search."); - int nByteExtralSharedMemory = nByteReservedSharedMemoryPerBlock + nByteStaticSharedMemory; - int nBlock = nMaxActiveBlock; - int nVPart = kMaxVPartStage1 + 1; - for (; nBlock > 0 && nVPart > kMaxVPartStage1; --nBlock) - { - int nByteDynamicSharedMemoryStage1 = nByteMaxSharedMemoryPerSM / nBlock - nByteExtralSharedMemory; - nByteDynamicSharedMemoryStage1 -= nByteDynamicSharedMemoryStage1 % sizeof(T); - nVPart = ceilDiv(sizeof(T) * vocabSize, nByteDynamicSharedMemoryStage1); - } - TLLM_CHECK_WITH_INFO(nBlock >= 0, "No enough active blocks for Beam Search stage 1 kernel."); - - int const nByteDynamicSharedMemoryStage1 = sizeof(T) * ceilDiv(vocabSize, nVPart); - this->mVPart = nVPart; - this->mByteSharedMemoryStage1 = nByteDynamicSharedMemoryStage1; // Only dynamic shared memory - - // Stage 2 - TLLM_CHECK_WITH_INFO(batchSize * beamWidth * paddedBeamWidth < (1 << 21), - "max_batch_size or max_beam_width of TRT-LLM engine is too large for Beam search, try to decrease the " - "parameters while building."); - size_t const nByteDynamicSharedMemoryStage2 = common::roundUp( - sizeof(float) * nVPart * (paddedBeamWidth * 4) + sizeof(cub::KeyValuePair) * paddedBeamWidth * 2, - 4); - switch (paddedBeamWidth) - { - case 1: GET_INFO_STAGE2(1); - case 2: GET_INFO_STAGE2(2); - case 4: GET_INFO_STAGE2(4); - case 8: GET_INFO_STAGE2(8); - default: break; - } - nByteStaticSharedMemory = attr.sharedSizeBytes; - nByteMaxDynamicSharedMemoryPerBlock = nByteMaxSharedMemoryPerBlock - nByteStaticSharedMemory; - nByteExtralSharedMemory = nByteReservedSharedMemoryPerBlock + nByteStaticSharedMemory; - bool const bUseGlobalMemoryStage2 = (nByteDynamicSharedMemoryStage2 > nByteMaxDynamicSharedMemoryPerBlock); - - // Stage 3 - // Keep top 2K candidates in case of k candidates finishes in one iteration - size_t const nByteDynamicSharedMemoryStage3 - = common::roundUp(sizeof(T) * paddedBeamWidth * paddedBeamWidth * 2, 4); - switch (paddedBeamWidth) - { - case 1: GET_INFO_STAGE3(1, false); - case 2: GET_INFO_STAGE3(2, false); - case 4: GET_INFO_STAGE3(4, false); - case 8: GET_INFO_STAGE3(8, false); - } - nByteStaticSharedMemory = attr.sharedSizeBytes; - nByteMaxDynamicSharedMemoryPerBlock = nByteMaxSharedMemoryPerBlock - nByteStaticSharedMemory; - nByteExtralSharedMemory = nByteReservedSharedMemoryPerBlock + nByteStaticSharedMemory; - bool const bUseGlobalMemoryStage3 = (nByteDynamicSharedMemoryStage3 > nByteMaxDynamicSharedMemoryPerBlock); - this->mByteSharedMemoryStage3 = nByteStaticSharedMemory; // Only static shared memory - - // Compute workspace size, see `beamSearchKernelsTemplate.h` for detailed information - // |<----- Workspace ----->| - // |<- A ->|<- B ->|<- C ->| - // |<---- D ---->| - // A for data exchange between stage 2 and 3 - // B for data exchange between stage 1 and 2, can be reuse for stage 3 - // C for stage 2 if `bUseGlobalMemoryStage2 == true`, can be reuse for stage 3 - // D for stage 3 if `bUseGlobalMemoryStage3 == true` - size_t const nByteA = common::roundUp(sizeof(T) * batchSize * paddedBeamWidth * paddedBeamWidth * 4, 4); - size_t const nByteB - = common::roundUp(sizeof(T) * batchSize * paddedBeamWidth * kMaxVPartStage1 * paddedBeamWidth * 4, 4); - size_t const nByteC = (bUseGlobalMemoryStage2) ? nByteDynamicSharedMemoryStage2 : 0; - size_t const nByteD = (bUseGlobalMemoryStage3) ? nByteDynamicSharedMemoryStage3 : 0; - this->mWorkspaceSize = nByteA + std::max(nByteB + nByteC, nByteD); - } - else // V2 workflow for large beam width or VBWS - { - this->mV2 = true; - switch (paddedBeamWidth) - { - case 1: GET_INFO_STAGE3(1, true); - case 2: GET_INFO_STAGE3(2, true); - case 4: GET_INFO_STAGE3(4, true); - case 8: GET_INFO_STAGE3(8, true); - case 16: GET_INFO_STAGE3(16, true); - case 32: GET_INFO_STAGE3(32, true); - case 64: GET_INFO_STAGE3(64, true); - case 128: GET_INFO_STAGE3(128, true); - case 256: GET_INFO_STAGE3(256, true); - case 512: GET_INFO_STAGE3(512, true); - case 1024: GET_INFO_STAGE3(1024, true); - } - this->mByteSharedMemoryStage3 = attr.sharedSizeBytes; // Only static shared memory - - // Compute shared memory size for stage 3 - // Compute workspace size, see `beamSearchKernelsTemplate.h` for detailed information - // |<----------------------------------------- Workspace ------------------------------------------>| - // |<- Stage2Ids ->|<- Stage2LogProbs ->|<- Stage1Ids ->|<- Stage1LogProbs ->|<---- Stage1TopK ---->| - // |<- stage2TopK ->| - // |<------------------ Stage3 ------------------>| - SizeType32 const batchSize{mDecoderDomain.getBatchSize()}; - SizeType32 const beamWidth{mDecoderDomain.getBeamWidth()}; - SizeType32 const vocabSize{mDecoderDomain.getVocabSize()}; - SizeType32 const paddedBeamWidth{padToNextPowerOfTwo(beamWidth)}; - size_t const nByteStage1LogProbs = roundUp(sizeof(T) * batchSize * paddedBeamWidth * paddedBeamWidth * 2, 4); - size_t const nByteStage1Ids = roundUp(sizeof(int) * batchSize * paddedBeamWidth * paddedBeamWidth * 2, 4); - size_t const nByteStage2LogProbs = roundUp(sizeof(T) * batchSize * paddedBeamWidth * 2, 4); - size_t const nByteStage2Ids = roundUp(sizeof(int) * batchSize * paddedBeamWidth * 2, 4); - size_t const nByteStage1TopK - = invokeComputeTopkLastDimWorkspaceSize(batchSize * beamWidth, vocabSize, paddedBeamWidth * 2, true); - size_t const nByteStage2TopK = invokeComputeTopkLastDimWorkspaceSize( - batchSize, paddedBeamWidth * paddedBeamWidth * 2, beamWidth * 2, true); - size_t const nByteStage3 = sizeof(T) * beamWidth * beamWidth * 2; - this->mWorkspaceSize = nByteStage2LogProbs + nByteStage2Ids - + max(nByteStage1LogProbs + nByteStage1Ids + max(nByteStage1TopK, nByteStage2TopK), nByteStage3); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t BeamSearchLayer::getWorkspaceSize() const noexcept -{ - return mWorkspaceSize; -} - -template -void BeamSearchLayer::setup(SizeType32 const batchSize, SizeType32 const beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(BeamSearchLayer_setup); - - SizeType32 const maxBamWidth{mDecoderDomain.getBeamWidth()}; - TLLM_CHECK_WITH_INFO(beamWidth <= maxBamWidth, "Beam width is larger than the constructed for (%d > %d).", - int(beamWidth), int(maxBamWidth)); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - auto constexpr fltMax = std::numeric_limits::max(); - auto constexpr fltMin = std::numeric_limits::lowest(); - auto constexpr fltEpsilon = std::numeric_limits::epsilon(); - auto constexpr int32Max = std::numeric_limits::max(); - FillBuffers const fillBuffers{batchSize, mDecoderDomain.getBatchSize(), mBufferManager}; - fillBuffers(setupParams->beamSearchDiversityRate, DefaultDecodingParams::getBeamSearchDiversity(), - mBeamSearchDiversityRateHost, mBeamSearchDiversityRateDevice, batchSlots, std::make_pair(-fltEpsilon, fltMax), - "diversity rate"); - fillBuffers(setupParams->lengthPenalty, DefaultDecodingParams::getLengthPenalty(), mLengthPenaltyHost, - mLengthPenaltyDevice, batchSlots, std::make_pair(fltMin, fltMax), "length penalty"); - fillBuffers(setupParams->earlyStopping, DefaultDecodingParams::getEarlyStopping(), mEarlyStoppingHost, - mEarlyStoppingDevice, batchSlots, std::make_pair(-fltEpsilon, int32Max), "early stopping"); - - if (this->mVBWS) - { - fillBuffers(setupParams->beamWidthArray, DefaultDecodingParams::getBeamWidthArray(), mBeamWidthArrayHost, - mBeamWidthArrayDevice, batchSlots, std::make_pair(-fltEpsilon, kMaxBeamWidth), "beam width array"); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void BeamSearchLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(BeamSearchLayer_forwardAsync); - - auto ip = std::dynamic_pointer_cast(baseInputs); - auto op = std::dynamic_pointer_cast(baseOutputs); - auto const localDecoderDomain = getLocalDecoderDomain(ip, mDecoderDomain); - - TLLM_CHECK_WITH_INFO(localDecoderDomain.getBeamWidth() > 1, "Use beamWidth <= 1 (%d <= 1) in Beam Search mode", - localDecoderDomain.getBeamWidth()); - TLLM_CHECK_WITH_INFO(ip->srcCacheIndirection.has_value(), "srcCacheIndirection is mandatory in beam search."); - TLLM_CHECK_WITH_INFO(op->parentIds.has_value(), "parentIds tensor is mandatory in beam search."); - TLLM_CHECK_WITH_INFO(op->finished.has_value(), "finished tensor is mandatory in beam search."); - TLLM_CHECK_WITH_INFO(op->cumLogProbs.has_value(), "cumLogProbs tensor is mandatory in beam search."); - TLLM_CHECK_WITH_INFO(op->beamHypotheses, "Output BeamHypotheses is not set."); - TLLM_CHECK_WITH_INFO(bufferCastOrNull(*op->sequenceLength) != nullptr || mLengthPenaltyDevice == nullptr, - "Current sequence lengths must be set for length penalty computation."); - TLLM_CHECK_WITH_INFO(ip->ite == 0, "Pipeline Parallelism is not supported yet!"); - - BeamHypotheses bh; - // bh's members not used in this function: outputIds, logProbs, outputIdsUnfinish, parentIdsUnfinish - bh.bVBWS = this->mVBWS; - // outputIds retains its full maxBatchSize allocation; outputIdsPtr is sliced to the active - // batch size in DynamicDecodeLayer::prepareIdsPtrs (ITensor::slice(mOutputIdsPtrDevice, 0, batchSize)) - // and must not be used as a stride for the [MSL, maxBatchSize, BM]-shaped logProbsTiled buffer. - bh.nMaxBatchSize = static_cast(op->outputIds->getDimension<0>()); - bh.nBatchSize = ip->localBatchSize; - bh.nBeamWidth = op->outputIds->getDimension<1>(); - bh.nMaxSeqLen = op->outputIds->getDimension<2>(); - bh.nVocabSize = mDecoderDomain.getVocabSizePadded(); - bh.nVPart = this->mVPart; - bh.nByteMaxSharedMemoryPerBlock = this->mByteMaxSharedMemoryPerBlock; - bh.nByteSharedMemoryStage1 = this->mByteSharedMemoryStage1; - bh.nByteSharedMemoryStage3 = this->mByteSharedMemoryStage3; - bh.diversityRates = bufferCast(*mBeamSearchDiversityRateDevice); - bh.lengthPenalties = bufferCast(*mLengthPenaltyDevice); - bh.earlyStoppings = bufferCast(*mEarlyStoppingDevice); - - if (this->mVBWS) - { - bh.beamWidthArraysHost = bufferCast(*mBeamWidthArrayHost); - bh.nBeamWidthInHost = bufferCast(*mBeamWidthIn); - bh.nBeamWidthOutHost = bufferCast(*mBeamWidthOut); - int const* batchSlotsHost = bufferCast(*ip->batchSlots); - for (int i = 0; i < ip->localBatchSize; ++i) - { - auto const slot = batchSlotsHost[i]; - auto const step = ip->beamSearchSteps.value()[slot]; - // Clamp `step` to [0, kMaxBeamWidthArrayLength - 1], and set `indexInput=0` when step = 0 or 1 - auto const indexOutput = std::min(step, static_cast(kMaxBeamWidthArrayLength) - 1); - auto const indexInput = std::max(indexOutput - 1, 0); - bh.nBeamWidthInHost[i] = bh.beamWidthArraysHost[slot * kMaxBeamWidthArrayLength + indexInput]; - bh.nBeamWidthOutHost[i] = bh.beamWidthArraysHost[slot * kMaxBeamWidthArrayLength + indexOutput]; - } - // At present, all requests of a batch must have the same beam width in one generation step (or they will not - // be batched together). So, the beam width of the first request is taken here to reshape the buffer. - // Corresponding changes must be done if Diverse-Beam-Width-Search (DBWS, requests with diverse beam width in - // a batch in one generation step) is supported in the future. - op->beamWidth = bh.nBeamWidthOutHost[0]; - } - else - { - op->beamWidth = bh.nBeamWidth; - } - - bh.inputLengths = bufferCast(*ip->inputLengths.value()); - bh.endIds = bufferCast(*ip->endIds); - bh.batchSlots = workspace->getDeviceBatchSlotsPtr(); // Device copy of `ip->batchSlots` - bh.logProbsTiled = bufferCastOrNull(op->outputLogProbsTiled); - bh.sequenceLengths = bufferCast(*op->sequenceLength.value()); - bh.cumLogProbs = bufferCast(*op->cumLogProbs.value()); - bh.outputIdsCBA = op->beamHypotheses->outputIdsCBA; - bh.logProbsCBA = op->beamHypotheses->logProbsCBA; - bh.sequenceLengthsCBA = op->beamHypotheses->sequenceLengthsCBA; - bh.cumLogProbsCBA = op->beamHypotheses->cumLogProbsCBA; - bh.normedScoresCBA = op->beamHypotheses->normedScoresCBA; - bh.numBeamsCBA = op->beamHypotheses->numBeamsCBA; - bh.minNormedScoresCBA = op->beamHypotheses->minNormedScoresCBA; - bh.batchDones = op->beamHypotheses->batchDones; - bh.finished = reinterpret_cast(bufferCast(*op->finished.value())); - bh.outputIdsPtr = bufferCast(*op->outputIdsPtr); - bh.parentIdsPtr = bufferCast(*op->parentIdsPtr); - - T const* logProbs = bufferCast(*workspace->getDeviceRuntimeLogits()); - T const* bias = static_cast(nullptr); - TLLM_CHECK_WITH_INFO(getWorkspaceSize() >= 2 * bh.nBatchSize * bh.nBeamWidth * bh.nBeamWidth * 2, - "Workspace size (%lu) is not enough for topk softmax required (%lu).", (uint64_t) getWorkspaceSize(), - (uint64_t) (2 * bh.nBatchSize * bh.nBeamWidth * bh.nBeamWidth * 2)); - - if (this->mV2 || this->mVBWS) - { - invokeTopkBeamSearch(logProbs, bias, workspace->getRawWorkspaceDevicePtr(), bh, getStream()); - } - else - { - invokeTopkBeamSearch(logProbs, bias, workspace->getRawWorkspaceDevicePtr(), bh, getStream()); - } - - int* tgtCI = bufferCast(*op->tgtCacheIndirection); - int* srcCI = bufferCast(*ip->srcCacheIndirection.value()); - invokeUpdateCacheIndirection(tgtCI, srcCI, bh, ip->maxAttentionWindow, ip->sinkTokenLength, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class BeamSearchLayer; -template class BeamSearchLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/beamSearchLayer.h b/cpp/tensorrt_llm/layers/beamSearchLayer.h deleted file mode 100644 index bff95a4f691b..000000000000 --- a/cpp/tensorrt_llm/layers/beamSearchLayer.h +++ /dev/null @@ -1,75 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" - -namespace tensorrt_llm::layers -{ - -template -class BeamSearchLayer : public BaseLayer -{ - using Base = BaseLayer; - -public: - BeamSearchLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager); - - // Functions called before input data arrives - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - - // Functions called after input data arrives - void setup(runtime::SizeType32 const batchSize, runtime::SizeType32 const beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - -private: - // Functions called before input data arrives - void allocateBuffer(); - void configureBeamSearchLayer(); - -private: - using Base::mDecoderDomain; - - size_t mByteMaxSharedMemoryPerBlock{0}; // Device information - size_t mByteSharedMemoryStage1{0}; // Max dynamic shashared memoryn stage 1 kernel, useless in V2 - size_t mByteSharedMemoryStage3{0}; // Max static shared memory in stage 3 kernel - size_t mVPart{0}; // Count of parts the beamed-logProbs will be divided into, useless in V2 - size_t mWorkspaceSize{0}; // Total workspace size for Beam Search kernels - bool mV2{false}; // Whether to use V2 Beam Search kernels - bool mVBWS{false}; // Whether to use Variable-Beam-Width-Search - - TensorPtr mBeamSearchDiversityRateHost; // [batchSize] cpu - TensorPtr mBeamSearchDiversityRateDevice; // [batchSize] gpu - TensorPtr mLengthPenaltyHost; // [batchSize] cpu - TensorPtr mLengthPenaltyDevice; // [batchSize] gpu - TensorPtr mEarlyStoppingHost; // [batchSize] cpu - TensorPtr mEarlyStoppingDevice; // [batchSize] gpu - TensorPtr mBeamWidthArrayHost; // [batchSize, kMaxBeamWidthArrayLength] cpu - TensorPtr mBeamWidthArrayDevice; // [batchSize, kMaxBeamWidthArrayLength] gpu - TensorPtr mBeamWidthIn; // [batchSize] cpu, the beamWidth of last forward computation - TensorPtr mBeamWidthOut; // [batchSize] cpu, the beamWidth of next forward computation -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/decodingLayer.cpp b/cpp/tensorrt_llm/layers/decodingLayer.cpp deleted file mode 100644 index 9d69d208a802..000000000000 --- a/cpp/tensorrt_llm/layers/decodingLayer.cpp +++ /dev/null @@ -1,266 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "decodingLayer.h" -#include "tensorrt_llm/layers/beamSearchLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/layers/eagleDecodingLayer.h" -#include "tensorrt_llm/layers/explicitDraftTokensLayer.h" -#include "tensorrt_llm/layers/externalDraftTokensLayer.h" -#include "tensorrt_llm/layers/layerUtils.h" -#include "tensorrt_llm/layers/lookaheadDecodingLayer.h" -#include "tensorrt_llm/layers/medusaDecodingLayer.h" -#include "tensorrt_llm/layers/samplingLayer.h" - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -DecodingLayer::DecodingLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (mDecodingMode.isTopKorTopP()) - { - mDecodingLayer = std::make_unique>(mDecodingMode, decoderDomain, mBufferManager); - } - else if (mDecodingMode.isBeamSearch()) - { - mDecodingLayer = std::make_unique>(mDecodingMode, decoderDomain, mBufferManager); - } - else if (mDecodingMode.isMedusa()) - { - mDecodingLayer = std::make_unique>(decoderDomain, mBufferManager); - } - else if (mDecodingMode.isLookahead()) - { - mDecodingLayer = std::make_unique>(mDecoderDomain, mBufferManager); - } - else if (mDecodingMode.isExplicitDraftTokens()) - { - mDecodingLayer = std::make_unique>(decoderDomain, mBufferManager); - } - else if (mDecodingMode.isExternalDraftTokens()) - { - mDecodingLayer = std::make_unique>(mDecodingMode, decoderDomain, mBufferManager); - } - else if (mDecodingMode.isEagle()) - { - mDecodingLayer = std::make_unique>(decoderDomain, mBufferManager); - } - else - { - TLLM_CHECK_WITH_INFO(false, - "Decoding mode is none of the supported {TopK, TopP, TopKTopP, BeamSearch, Medusa, Lookahead, " - "ExplicitDraftTokens, ExternalDraftTokens, Eagle}"); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DecodingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - TLLM_CHECK_WITH_INFO(setupParams->decodingParams, "decodingParams for setup is not set"); - - if (mDecodingMode.isBeamSearch()) - { - TLLM_CHECK_WITH_INFO( - beamWidth > 1, "Decoding mode is %s, but beamWidth <= 1 (%d <= 1)", mDecodingMode.getName(), beamWidth); - } - else if (mDecodingMode.isTopKorTopP() || mDecodingMode.isMedusa() || mDecodingMode.isLookahead() - || mDecodingMode.isExplicitDraftTokens() || mDecodingMode.isExternalDraftTokens() || mDecodingMode.isEagle()) - { - TLLM_CHECK_WITH_INFO( - beamWidth == 1, "Decoding mode is %s, but beamWidth != 1 (%d != 1)", mDecodingMode.getName(), beamWidth); - } - else - { - TLLM_CHECK_WITH_INFO(false, - "Decoding mode is none of the supported {TopK, TopP, TopKTopP, BeamSearch, Medusa, Lookahead, " - "ExplicitDraftTokens, ExternalDraftTokens, Eagle}"); - } - - mDecodingLayer->setup(batchSize, beamWidth, batchSlots, setupParams->decodingParams, workspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DecodingLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto [outputParams, inputParams] = prepareParams(baseOutputs, baseInputs); - mDecodingLayer->forwardAsync(outputParams, inputParams, workspace); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DecodingLayer::forwardSync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto [outputParams, inputParams] = prepareParams(baseOutputs, baseInputs); - mDecodingLayer->forwardSync(outputParams, inputParams, workspace); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t DecodingLayer::getWorkspaceSize() const noexcept -{ - return mDecodingLayer->getWorkspaceSize(); -} - -template -std::tuple, std::shared_ptr> DecodingLayer::prepareParams( - std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto params = std::dynamic_pointer_cast(baseInputs); - - auto const localDecoderDomain = getLocalDecoderDomain(params, mDecoderDomain); - auto const& endIds = params->endIds; - - std::shared_ptr preparedOutputs; - std::shared_ptr preparedInputs; - - if (mDecodingMode.isBeamSearch()) - { - preparedInputs = baseInputs; - preparedOutputs = baseOutputs; - } - else if (mDecodingMode.isTopKorTopP()) - { - auto const ite = params->ite; - auto const step = params->step; - auto const localBatchSize = static_cast(params->localBatchSize); - - TLLM_CHECK_WITH_INFO(localDecoderDomain.getBeamWidth() == 1, - "Decoding mode is TopK and/or TopP, but beamWidth != 1 (%d != 1)", localDecoderDomain.getBeamWidth()); - - // In sampling, we have supported batch sampling. So, we always compute all - // sentences once. - TensorConstPtr logitsSlice = ITensor::slice(*params->logits, 0, localBatchSize); - TensorConstPtr endIdSlice = ITensor::slice(endIds, 0, localBatchSize); - auto decodeInputs = std::make_shared(endIdSlice, params->batchSlots, step, ite, localBatchSize); - - decodeInputs->finished = params->finished; - - decodeInputs->logits = logitsSlice; - - if (params->inputLengths) - { - auto& inputLengths = params->inputLengths.value(); - decodeInputs->inputLengths = ITensor::slice(inputLengths, 0, localBatchSize); - } - preparedInputs = decodeInputs; - preparedOutputs = baseOutputs; - } - else if (mDecodingMode.isMedusa()) - { - TLLM_CHECK_WITH_INFO(localDecoderDomain.getBeamWidth() == 1, - "Decoding mode is Medusa, but beamWidth != 1 (%d != 1)", localDecoderDomain.getBeamWidth()); - - preparedInputs = baseInputs; - preparedOutputs = baseOutputs; - } - else if (mDecodingMode.isLookahead()) - { - preparedInputs = baseInputs; - preparedOutputs = baseOutputs; - } - else if (mDecodingMode.isExplicitDraftTokens()) - { - preparedInputs = baseInputs; - preparedOutputs = baseOutputs; - } - else if (mDecodingMode.isExternalDraftTokens()) - { - auto externalDraftTokenParams = std::dynamic_pointer_cast(baseInputs); - auto const ite = externalDraftTokenParams->ite; - auto const step = externalDraftTokenParams->step; - auto const localBatchSize = static_cast(externalDraftTokenParams->localBatchSize); - - TLLM_CHECK_WITH_INFO(localDecoderDomain.getBeamWidth() == 1, - "Decoding mode is TopK and/or TopP, but beamWidth != 1 (%d != 1)", localDecoderDomain.getBeamWidth()); - - // Compute all sentences once since batch-sampling is supported - TensorConstPtr logitsSlice = ITensor::slice(*externalDraftTokenParams->logits, 0, localBatchSize); - TensorConstPtr endIdSlice = ITensor::slice(endIds, 0, localBatchSize); - auto decodeInputs = std::make_shared( - endIdSlice, externalDraftTokenParams->batchSlots, step, ite, localBatchSize); - - decodeInputs->finished = externalDraftTokenParams->finished; - - decodeInputs->logits = logitsSlice; - - if (externalDraftTokenParams->inputLengths) - { - auto& inputLengths = externalDraftTokenParams->inputLengths.value(); - decodeInputs->inputLengths = ITensor::slice(inputLengths, 0, localBatchSize); - } - decodeInputs->draftLogits = externalDraftTokenParams->draftLogits; - decodeInputs->draftProbs = externalDraftTokenParams->draftProbs; - decodeInputs->targetProbs = externalDraftTokenParams->targetProbs; - decodeInputs->numDraftTokens = externalDraftTokenParams->numDraftTokens; - decodeInputs->numDraftTokensHost = externalDraftTokenParams->numDraftTokensHost; - decodeInputs->draftTokenIds = externalDraftTokenParams->draftTokenIds; - decodeInputs->constantThreshold = externalDraftTokenParams->constantThreshold; - decodeInputs->useRandomAcceptanceThreshold = externalDraftTokenParams->useRandomAcceptanceThreshold; - decodeInputs->step = externalDraftTokenParams->step; - decodeInputs->useDraftLogits = externalDraftTokenParams->useDraftLogits; - decodeInputs->useDraftLogitsHost = externalDraftTokenParams->useDraftLogitsHost; - - preparedInputs = decodeInputs; - preparedOutputs = baseOutputs; - } - else if (mDecodingMode.isEagle()) - { - preparedInputs = baseInputs; - preparedOutputs = baseOutputs; - } - else - { - TLLM_CHECK_WITH_INFO(false, - "Decoding mode is none of the supported {TopK, TopP, TopKTopP, BeamSearch, Medusa, Lookahead, " - "ExplicitDraftTokens, ExternalDraftTokens, Eagle}"); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {preparedOutputs, preparedInputs}; -} - -template class DecodingLayer; -template class DecodingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/decodingLayer.h b/cpp/tensorrt_llm/layers/decodingLayer.h deleted file mode 100644 index 60780851f977..000000000000 --- a/cpp/tensorrt_llm/layers/decodingLayer.h +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Layer performs token decoding using sampling (beamWidth=1), beam search (beamWidth>1) or Medusa. -template -class DecodingLayer : public BaseLayer -{ -public: - DecodingLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - //! \brief Calls single SamplingLayer::forwardAsync or MedusaDecodingLayer::forwardAsync in batched mode - //! or runs BeamSearchLayer::forwardAsync in the loop for each request. - //! Modifies outputs->logits in-place. - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! \brief Calls forwardSync of configured decoding layer. - void forwardSync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - [[nodiscard]] std::tuple, std::shared_ptr> prepareParams( - std::shared_ptr const& outputs, std::shared_ptr const& inputs) const; - -private: - using BaseLayer::mDecoderDomain; - - executor::DecodingMode mDecodingMode; - - std::unique_ptr mDecodingLayer; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/decodingParams.h b/cpp/tensorrt_llm/layers/decodingParams.h deleted file mode 100644 index 76c5cedd637b..000000000000 --- a/cpp/tensorrt_llm/layers/decodingParams.h +++ /dev/null @@ -1,689 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include -#include - -#include -#include -#include - -namespace tensorrt_llm::layers -{ - -using TensorPtr = runtime::ITensor::SharedPtr; -using TensorConstPtr = runtime::ITensor::SharedConstPtr; -using BufferPtr = runtime::IBuffer::SharedPtr; -using BufferConstPtr = runtime::IBuffer::SharedConstPtr; -template -using OptVec = std::optional>; - -//! -//! \brief In a DecodingLayer's life cycle, it is constructed once; -//! `setup` repeatedly, but once per request; `forward*` repeatedly, many times per request. -//! A possible sequence would be, construct(maxBatchSize) -> setup({1,3}) -> forward({1, 3}) -//! -> forward({1, 3}) -> setup({2,4}) -> forward({1, 3, 2, 4}) -> forward({1, 3, 2, 4}) -//! -> forward({1, 2, 4}), where {a,b} are batchSlots, and {3} ends at last step. -//! As a result there are three types of batches. -//! 1. `maxBatchSize` for each layers to reserve resources. -//! It is passed through class constructor, in DecoderDomain.getBatchSize(). -//! 2. `setupBatchSize` for setting up layers for a batch of new requests. -//! It is passed through `setup` method. -//! 3. `forwardBatchSize` for layers forwarding for a batch of existing active requests. -//! it is passed through `forwardAsync` and `forwardSync` methods. -//! `setup` and `forward` always provide `batchSlots` indexed by -//! local batch index ranging in [0, setupBatchSize) or [0, forwardBatchSize), -//! holding the global batch index ranging in [0, maxBatchSize). -//! In case of beam search, maxBatchSize = forwardBatchSize = 1. - -class DecoderDomain -{ -public: - DecoderDomain(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, runtime::SizeType32 vocabSize, - std::optional vocabSizePadded = std::nullopt, - std::shared_ptr speculativeDecodingModule = nullptr) - : mBatchSize(batchSize) - , mBeamWidth(beamWidth) - , mVocabSize(vocabSize) - , mVocabSizePadded(vocabSizePadded.value_or(vocabSize)) - , mSpeculativeDecodingModule(std::move(speculativeDecodingModule)) - { - } - - [[nodiscard]] runtime::SizeType32 getBatchSize() const - { - return mBatchSize; - } - - [[nodiscard]] runtime::SizeType32 getBeamWidth() const - { - return mBeamWidth; - } - - void setBeamWidth(runtime::SizeType32 beamWidth) - { - mBeamWidth = beamWidth; - } - - [[nodiscard]] runtime::SizeType32 getVocabSize() const - { - return mVocabSize; - } - - [[nodiscard]] runtime::SizeType32 getVocabSizePadded() const - { - return mVocabSizePadded; - } - - [[nodiscard]] runtime::SizeType32 getMaxDecodingTokens() const - { - return mSpeculativeDecodingModule ? mSpeculativeDecodingModule->getMaxDecodingTokens() : 1; - } - - [[nodiscard]] std::shared_ptr getSpeculativeDecodingModule() const - { - TLLM_CHECK_WITH_INFO(mSpeculativeDecodingModule, "Speculative decoding module is not set to decoder domain"); - return mSpeculativeDecodingModule; - } - - [[nodiscard]] std::shared_ptr getSpeculativeDecodingModulePtr() const - { - return mSpeculativeDecodingModule; - } - -private: - runtime::SizeType32 mBatchSize; - runtime::SizeType32 mBeamWidth; - runtime::SizeType32 mVocabSize; - runtime::SizeType32 mVocabSizePadded; - std::shared_ptr mSpeculativeDecodingModule; -}; - -class BaseSetupParams -{ -public: - virtual ~BaseSetupParams() = default; -}; - -// Penalty layer -class PenaltySetupParams : public BaseSetupParams -{ -public: - OptVec temperature; // [1] or [setupBatchSize] - OptVec minLength; // [1] or [setupBatchSize] - OptVec repetitionPenalty; // [1] or [setupBatchSize] - OptVec presencePenalty; // [1] or [setupBatchSize] - OptVec frequencyPenalty; // [1] or [setupBatchSize] - OptVec promptIgnoreLength; // [1] or [setupBatchSize] -}; - -// Ban words layer -class BanWordsSetupParams : public BaseSetupParams -{ -public: - OptVec noRepeatNgramSize; // [1] or [setupBatchSize] -}; - -class DecodingSetupParams : public BaseSetupParams -{ -public: - virtual ~DecodingSetupParams() = default; - - OptVec randomSeed; // [1] or [setupBatchSize] - OptVec outputLogProbs; // [setupBatchSize] - OptVec cumLogProbs; // [setupBatchSize] -}; - -class SamplingSetupParams : public DecodingSetupParams -{ -public: - // baseSamplingLayer - OptVec runtimeTopK; // [1] or [setupBatchSize] - OptVec runtimeTopP; // [1] or [setupBatchSize] - OptVec runtimeMinP; // [1] or [setupBatchSize] - - // topPSamplingLayer - OptVec topPDecay; // [setupBatchSize], between [0, 1] - OptVec topPMin; // [setupBatchSize], between [0, 1] - OptVec topPResetIds; // [setupBatchSize] - std::optional normalizeLogProbs; -}; - -class BeamSearchSetupParams : public DecodingSetupParams -{ -public: - // BeamSearchLayer - OptVec beamSearchDiversityRate; // [setupBatchSize] - OptVec lengthPenalty; // [setupBatchSize] - OptVec earlyStopping; // [setupBatchSize] - OptVec> beamWidthArray; // [setupBatchSize, nMaxBeamWidthArray] - bool hasDiffRuntimeArgs{false}; -}; - -class MedusaSetupParams : public DecodingSetupParams -{ -public: - // Medusa params - OptVec runtimeTopK; // [setupBatchSize] - OptVec> runtimeHeadsTopK; // [setupBatchSize, maxMedusaHeads] -}; - -class ExplicitDraftTokensSetupParams : public DecodingSetupParams -{ -public: - OptVec temperature; // [setupBatchSize] - // Hack to init some data for the context phase in the setup. - TensorPtr randomDataSample; // [maxBatchSize], on gpu - TensorPtr temperatures; // [maxBatchSize], on gpu - tensorrt_llm::DataType dtype; // [1] -}; - -class EagleSetupParams : public DecodingSetupParams -{ -public: - OptVec temperature; // [setupBatchSize] - // Hack to init some data for the context phase in the setup. - TensorPtr randomDataSample; // [maxBatchSize], on gpu - TensorPtr temperatures; // [maxBatchSize], on gpu - tensorrt_llm::DataType dtype; // [1] -}; - -class DynamicDecodeSetupParams : public BaseSetupParams -{ -public: - std::shared_ptr penaltyParams; - std::shared_ptr banWordsParams; - std::shared_ptr decodingParams; -}; - -struct LookaheadSetupParams : public DecodingSetupParams -{ - using TensorPtr = runtime::ITensor::SharedPtr; - - std::vector prompt; // [batchSize][maxSeqLen], on cpu - std::vector algoConfigs; // [1] or [batchSize] - - //! see class LookaheadDecodingOutputs - TensorPtr generationLengths; // [maxBatchSize], on gpu - TensorPtr positionOffsets; // [maxBatchSize, maxDecodingTokens], on gpu - TensorPtr attentionPackedMasks; // [maxBatchSize, maxDecodingTokens], on gpu -}; - -class ExternalDraftTokensSetupParams : public DecodingSetupParams -{ -public: - OptVec runtimeTopK; // [1] or [setupBatchSize] - OptVec runtimeTopP; // [1] or [setupBatchSize] -}; - -class BaseDecodingInputs -{ -public: - BaseDecodingInputs(runtime::SizeType32 localBatchSize) - : localBatchSize(localBatchSize) - { - } - - virtual ~BaseDecodingInputs() = default; - - runtime::SizeType32 localBatchSize; -}; - -// Ban words inputs -class BanWordsDecodingInputs : public BaseDecodingInputs -{ -public: - BanWordsDecodingInputs(runtime::SizeType32 localBatchSize) - : BaseDecodingInputs(localBatchSize) - { - } - - runtime::SizeType32 maxBadWordsLen{0}; - std::optional badWordsPtr; // [maxBatchSize][2, bad_words_length], on gpu - std::optional badWordsLengths; // [maxBatchSize], on gpu -}; - -// Stop criteria inputs -class StopCriteriaDecodingInputs : public BaseDecodingInputs -{ -public: - StopCriteriaDecodingInputs(runtime::SizeType32 localBatchSize) - : BaseDecodingInputs(localBatchSize) - { - } - - runtime::SizeType32 maxStopWordsLen{0}; - std::optional sequenceLimitLength; // [maxBatchSize], on gpu - std::optional stopWordsPtr; // [maxBatchSize][2, stop_words_length], on pinned - std::optional stopWordsLengths; // [maxBatchSize], on pinned -}; - -class DecodingInputs : public BaseDecodingInputs -{ -public: - DecodingInputs(TensorConstPtr endIds, TensorConstPtr batchSlots, runtime::SizeType32 step = 0, - runtime::SizeType32 ite = 0, runtime::SizeType32 localBatchSize = 0, runtime::SizeType32 maxAttentionWindow = 0, - runtime::SizeType32 sinkTokenLength = 0) - : BaseDecodingInputs(localBatchSize) - , endIds{std::move(endIds)} - , step{step} - , ite{ite} - , maxAttentionWindow{maxAttentionWindow} - , sinkTokenLength{sinkTokenLength} - , batchSlots{std::move(batchSlots)} - { - } - - TensorConstPtr endIds; // [maxBatchSize] - - // used only for python runtime - runtime::SizeType32 step; - runtime::SizeType32 ite; - - // mandatory parameters - runtime::SizeType32 maxAttentionWindow; - runtime::SizeType32 sinkTokenLength; - - //! One of `logits` and `logitsVec` has to be set - //! DynamicDecodeLayer::forward checks for it - //! Need both of these fields to support legacy code during transition period to the batched decoder - std::optional logits; // [forwardBatchSize, beamWidth, vocabSizePadded], on gpu - OptVec logitsVec; // [forwardBatchSize][beamWidth, vocabSizePadded], on gpu - TensorConstPtr batchSlots; // [forwardBatchSize], on pinned - - // optional parameters - std::optional - srcCacheIndirection; // [forwardBatchSize, maxBeamWidth, maxSeqLen], on gpu, for Beam Search - std::optional embeddingBias; // [vocabSizePadded], on gpu - std::optional inputLengths; // [maxBatchSize, maxBeamWidth], on gpu - std::optional finished; // [maxBatchSize, maxBeamWidth] - std::optional curTokensPerStep; // [maxBatchSize], on gpu - std::shared_ptr banWordsInputs; - std::shared_ptr stopCriteriaInputs; - OptVec beamSearchSteps; // [forwardBatchSize], for Variable-Beam-Width-Search -}; - -class SamplingInputs : public DecodingInputs -{ -public: - explicit SamplingInputs(TensorConstPtr endIds, TensorConstPtr batchSlots, runtime::SizeType32 step, - runtime::SizeType32 ite, runtime::SizeType32 localBatchSize) - : DecodingInputs{std::move(endIds), std::move(batchSlots), step, ite, localBatchSize} - { - } - - //! optional parameters - curandState_t* curandStates{}; // [localBatchSize] - - //! Flag to mark that logits tensor contains probabilities - bool probsComputed{}; -}; - -class ExternalDraftTokensInputs : public DecodingInputs -{ -public: - explicit ExternalDraftTokensInputs(TensorConstPtr endIds, TensorConstPtr batchSlots, runtime::SizeType32 step, - runtime::SizeType32 ite, runtime::SizeType32 localBatchSize) - : DecodingInputs{std::move(endIds), std::move(batchSlots), step, ite, localBatchSize} - { - } - - TensorPtr draftLogits; - TensorPtr draftProbs; - TensorPtr targetProbs; - TensorPtr numDraftTokens; - TensorPtr numDraftTokensHost; - TensorPtr draftTokenIds; - TensorPtr useDraftLogits; - TensorPtr useDraftLogitsHost; - - runtime::SizeType32 step{}; - float constantThreshold{}; - bool useRandomAcceptanceThreshold{}; - - //! optional parameters - curandState_t* curandStates{}; // [localBatchSize] - - //! Flag to mark that logits tensor contains probabilities - bool probsComputed{}; -}; - -// Medusa inputs -class MedusaDecodingInputs : public DecodingInputs -{ -public: - explicit MedusaDecodingInputs(TensorConstPtr endIds, TensorConstPtr batchSlots, runtime::SizeType32 localBatchSize) - : DecodingInputs(std::move(endIds), std::move(batchSlots), 0, 0, localBatchSize) - { - } - - TensorConstPtr targetTokensPerStep; // [maxBatchSize], on gpu - TensorConstPtr paths; // [maxBatchSize, maxPathLen, maxPathLen], on gpu - TensorConstPtr treeIds; // [maxBatchSize, maxDecodingTokens], on gpu - - // [maxBatchSize][maxDraftPathLen][maxDecodingTokens, vocabSizePadded], on gpu - std::vector> medusaLogits; -}; - -// Explicit draft tokens inputs -class ExplicitDraftTokensInputs : public DecodingInputs -{ -public: - explicit ExplicitDraftTokensInputs(TensorConstPtr endIds, TensorConstPtr batchSlots, runtime::SizeType32 batchSize) - : DecodingInputs(std::move(endIds), std::move(batchSlots), 0, 0, batchSize) - { - } - - //! Draft tokens for the next iteration. The first token in each path is the last accepted token at current - //! iteration. E.g. if forwardBatchSize == 1, maxNumPaths == 2, maxPathLen== 3, [[[0, 1, 2], [0, 1, 10]]] - TensorConstPtr nextDraftTokens; // [forwardBatchSize, maxNumPaths, maxPathLen], gpu - //! Compressed form of `nextDraftTokens`, where common prefixes and collapsed. - //! Using example above [0, 1, 2, 10] - TensorConstPtr nextFlatTokens; // [forwardBatchSize * maxDecodingTokens], gpu - //! Indices of draft tokens in the compressed `nextFlatTokens` for the next iteration. - //! Using example above, [[[0, 1, 2], [0, 1, 3]]] - TensorConstPtr nextDraftIndices; // [forwardBatchSize, maxNumPaths, maxPathLen], gpu - //! Probabilities of the next draft tokens. - TensorConstPtr nextDraftProbs; // [forwardBatchSize, maxNumPaths, maxDraftPathLen, vocabSize], gpu - //! Same as `nextDraftTokens`, but for current iteration. - //! Current accepted tokens obtained as `lastDraftTokens[bi][bestPathIndices[bi]][1:bestPathLengths[bi]]`. - TensorConstPtr lastDraftTokens; // [forwardBatchSize, maxNumPaths, maxPathLen], gpu - //! Same as `nextDraftIndices`, but for current iteration. - TensorConstPtr lastDraftIndices; // [forwardBatchSize, maxNumPaths, maxPathLen], gpu - //! Boolean attention masks. - //! maxDecodingTokens' = generationLengths.max() - TensorConstPtr masks; // [forwardBatchSize, maxDecodingTokens', maxDecodingTokens'], gpu - //! Relative to `positionIdsBase` position ids. Same as `nextFlatTokens` for next draft indices. - //! Using example above, [0, 1, 2, 3] - TensorConstPtr packedPosIds; // [forwardBatchSize * maxDecodingTokens], gpu - //! Lengths of the accepted paths for each request. It is 1 for context phase (Only 1 primary tokens is accepted). - TensorConstPtr bestPathLengths; // [forwardBatchSize], gpu - //! Indices of the accepted paths for each request. It is 0 for context phase. - TensorConstPtr bestPathIndices; // [forwardBatchSize], gpu - //! Number of the draft tokens for the next iteration. - TensorConstPtr generationLengths; // [forwardBatchSize], gpu - //! Baseline for the position ids. - TensorConstPtr positionIdsBase; // [forwardBatchSize], gpu - //! Generation length for the previous stage. - TensorConstPtr lastGenerationLengths; // [forwardBatchSize], gpu - //! Maximum number of generated tokens for the next step across whole batch - TensorConstPtr maxGenLengthDevice; // [1], on gpu - //! Address map to map from linear indices of the engine outputs to seqSlot. - //! It is not the same as batchSlots because it maps the ordered engine outputs to the respective seqSlot, - //! while batchSlots is just a a list of active seqSlots. - TensorConstPtr seqSlots; // [forwardBatchSize], on gpu -}; - -class LookaheadDecodingInputs : public DecodingInputs -{ -public: - explicit LookaheadDecodingInputs(TensorConstPtr endIds, TensorConstPtr batchSlots) - : DecodingInputs{std::move(endIds), std::move(batchSlots)} - { - } -}; - -// Explicit draft tokens inputs -class EagleInputs : public DecodingInputs -{ -public: - explicit EagleInputs(TensorConstPtr endIds, TensorConstPtr batchSlots, runtime::SizeType32 batchSize, - TensorConstPtr nextDraftTokens, TensorConstPtr nextDraftLens, TensorConstPtr nextDraftPaths, - TensorConstPtr lastDraftTokens, TensorConstPtr lastDraftLens, TensorConstPtr lastDraftPaths, - TensorConstPtr acceptedTokens, TensorConstPtr acceptedLens, TensorConstPtr acceptedPathIds, - TensorConstPtr chunkedContextNextTokens, TensorConstPtr seqSlots) - : DecodingInputs(std::move(endIds), std::move(batchSlots), 0, 0, batchSize) - , nextDraftTokens(nextDraftTokens) - , nextDraftLens(nextDraftLens) - , nextDraftPaths(nextDraftPaths) - , lastDraftTokens(lastDraftTokens) - , lastDraftLens(lastDraftLens) - , lastDraftPaths(lastDraftPaths) - , acceptedTokens(acceptedTokens) - , acceptedLens(acceptedLens) - , acceptedPathIds(acceptedPathIds) - , chunkedContextNextTokens(chunkedContextNextTokens) - , seqSlots(seqSlots) - { - } - - //! Draft tokens for the next iteration. - TensorConstPtr nextDraftTokens; // [forwardBatchSize, maxDecodingDraftTokens], gpu - //! Number of the draft tokens for the next iteration. - TensorConstPtr nextDraftLens; // [forwardBatchSize], gpu - //! Draft paths for the next iteration. - TensorConstPtr nextDraftPaths; // [forwardBatchSize, maxDecodingTokens, maxPathLen], gpu - //! Same as `nextDraftTokens`, but for current iteration. - TensorConstPtr lastDraftTokens; // [forwardBatchSize, maxNumPaths, maxPathLen], gpu - //! Number of the draft tokens input to the previous TRT iteration. - TensorConstPtr lastDraftLens; // [forwardBatchSize], gpu - //! Same as `nextDraftPaths`, but for current iteration. - TensorConstPtr lastDraftPaths; // [forwardBatchSize, maxDecodingTokens, maxPathLen], gpu - //! Lastly accepted tokens (including golden token). - TensorConstPtr acceptedTokens; // [forwardBatchSize, maxPathLen] - //! Number of accepted tokens (at least 1). - TensorConstPtr acceptedLens; // [forwardBatchSize] - //! Ids of the accepted path. - TensorConstPtr acceptedPathIds; // [forwardBatchSize] - //! Indicator whether the context request last chunk or not. - TensorConstPtr chunkedContextNextTokens; // [forwardBatchSize] - //! - TensorConstPtr seqSlots; // [forwardBatchSize], on gpu -}; - -class BaseDecodingOutputs -{ -public: - explicit BaseDecodingOutputs(TensorPtr outputIds) - : outputIds{std::move(outputIds)} - { - } - - virtual ~BaseDecodingOutputs() = default; - - //! Mandatory parameters - TensorPtr outputIds; // [maxBatchSize, maxSeqLen] - - //! Optional parameters - std::optional finished; // [maxBatchSize * maxBeamWidth], on pinned - std::optional sequenceLength; // [maxBatchSize * maxBeamWidth], on gpu - std::optional cumLogProbs; // [maxBatchSize * maxBeamWidth], on gpu, for Beam Search - //! NOTE: In the TRT backend, temperature is applied to logits in-place before sampling - //! (see decodingCommon.cu), so these log probs reflect the temperature-adjusted distribution. - //! This differs from the PyTorch backend's default behavior (LogprobMode.RAW), which computes - //! log probs from the raw model logits without temperature scaling. - std::optional outputLogProbs; // [maxBatchSize, maxBeamWidth, maxSeqLen], on gpu - std::optional parentIds; // [maxBatchSize, maxBeamWidth, maxSeqLen], on gpu, for Beam Search - - TensorPtr outputIdsPtr; // [maxBatchSize][maxBeamWidth, maxSeqLen], on gpu and outputIdsPtr[i], on gpu - TensorPtr outputIdsPtrHost; // [maxBatchSize][maxBeamWidth, maxSeqLen], on cpu but outputIdsPtrHost[i], on gpu - TensorPtr parentIdsPtr; // [maxBatchSize][maxBeamWidth, maxSeqLen], on cpu but parentIdsPtr[i], on gpu - TensorPtr newTokens; // [maxBatchSize, maxBeamWidth], on gpu, tokens predicted at current iteration. - - // optional parameters - std::optional numNewTokens; // [maxBatchSize], on pinned, number of tokens predicted at current iteration - std::optional finishedSum; // [1], on pinned - std::optional outputLogProbsTiled; // [maxSeqLen, maxBatchSize, maxBeamWidth], on gpu - - // Beam width might change in Variable-Beam-Width-Search mode. - // So the beam width is updated in beam search layer for the later layers. - runtime::SizeType32 beamWidth{1}; -}; - -class BeamSearchOutputs : public BaseDecodingOutputs -{ -public: - explicit BeamSearchOutputs(TensorPtr outputIds) - : BaseDecodingOutputs{std::move(outputIds)} - { - } - - TensorPtr tgtCacheIndirection; //[forwardBatchSize, maxBeamWidth, maxSeqLen], on gpu, the k/v cache index - - std::unique_ptr beamHypotheses; // Structure maintains variables of Beam Search -}; - -//! -//! \brief SpeculativeDecodingOutputs outputs. -//! -//! For one example sequence [a, b] [c] , where, [a, b, c] is the accepted sequence, -//! [c] is the last accepted token, and is the draft tokens from `nextDraftTokens` saved by last step. -//! [c]'s position id is known, only position ids for need to be provided in `nextDraftPosIds`. -//! LLM inputs {c, x, y, z} and generates {c', x', y', z'}. -//! -//! {c'} is always accepted and {x', z'} is supposed to be accepted. -//! The accepted tokens [c', x', z'] is saved in `outputIds` in-place, starting from `sequenceLength`. -//! The `acceptedLength` is 3, and the accepted draft tokens length is 2. -//! `sequenceLength` is also increased by `acceptedLength` in-place. -//! The pathsOffset is {0, 1, 3} for {c', x', z'}. -//! [] for accepted, <> for draft, {} for input/output. -//! -//! For a batchSlots {1, 3}, `numNewTokensCumSum` is an exclusive sum of `numNewTokens` over the batch, -//! the `numNewTokens` may be {3, 5}, `numNewTokensCumSum` is {0, 3, 8}. -//! -//! `nextDraftLengths` and `prevDraftLengths` are needed for methods that support if variable -//! draft length. `nextDraftLengths` must contain the number of draft tokens per request for the next iteration. -//! `prevDraftLengths` must contain the number of draft tokens used in the current iteraiton. -//! -//! `pathsOffsets` is needed for KV cache rewind. It contains the positions of the accepted draft tokens in the -//! flattened tensor of draft tokens. E.g. if for sequence {c, x, y, z} only `y` and `z` were accepted, -//! `pathsOffsets` contains [1, 2]. `pathsOffsets` is flattened tensor for whole batch. -//! -//! The order of `pathsOffsets` and `numNewTokensCumSum` must be aligned. Such that -//! `pathsOffset[numNewTokensCumSum[bi]:numNewTokensCumSum[bi+1]]` is the slice of offsets for `bi`th request. -//! Furthermore, the order of requests is important and must be aligned with sorted `RuntimeBuffers::seqSlots` -//! such that the request with smaller `seqSlot` stays earlier in the tensors. -//! However, this condition usually holds if method does not expect from the engine anything else, but logits. -class SpeculativeDecodingOutputs : public BaseDecodingOutputs -{ -public: - explicit SpeculativeDecodingOutputs(TensorPtr outputIds) - : BaseDecodingOutputs{std::move(outputIds)} - { - } - - //! Draft tokens for the next step - TensorPtr nextDraftTokens; // [maxBatchSize, maxDecodingDraftTokens] - //! Draft token position IDs - TensorPtr nextDraftPosIds; // [maxBatchSize, maxDecodingDraftTokens] - //! Prev step draft tokens lengths, should be filled only for variable draft length speculative decoding mode - TensorPtr prevDraftLengths; // [maxBatchSize] - //! Next step draft tokens lengths, should be filled only for variable draft length speculative decoding mode - TensorPtr nextDraftLengths; // [maxBatchSize] - //! Accumulative sum along batchSlots. - TensorPtr numNewTokensCumSum; // [maxBatchSize + 1] - TensorPtr pathsOffsets; // [maxBatchSize * maxPathLen] - TensorPtr packedMasks; // [maxBatchSize, maxDecodingTokens, divUp(maxDecodingTokens, 32)] -}; - -class LookaheadDecodingOutputs : public SpeculativeDecodingOutputs -{ - using TensorPtr = runtime::ITensor::SharedPtr; - -public: - explicit LookaheadDecodingOutputs(TensorPtr outputIds) - : SpeculativeDecodingOutputs{std::move(outputIds)} - { - } - - //! for TLLM engine input "spec_decoding_generation_lengths", indicating how many tokens to be generated. - //! currently, the 1st step of generation is 1, set at `setup`, others are maxDecodingTokens, set at `forward`. - TensorPtr generationLengths; // [maxBatchSize] - //! for TLLM engine input "spec_decoding_position_offsets", - //! indicating each token position offset base on the last golden token = 0. - //! ABCefgxyz--- // sequence tokens, ABCD: golden; efg, xyz: draft; ---: padding. - //! ***<0>123123--- // positionOffsets. - //! 012<3>456456--- // positionIds. - TensorPtr positionOffsets; // [maxBatchSize, maxDecodingTokens] - TensorPtr positionIds; // [maxBatchSize, maxDecodingTokens] -}; - -class ExplicitDraftTokensOutputs : public SpeculativeDecodingOutputs -{ -public: - explicit ExplicitDraftTokensOutputs(TensorPtr outputIds) - : SpeculativeDecodingOutputs{std::move(outputIds)} - { - } - - //! Draft tokens for the next iteration. The first token in each path is the last accepted token at current - //! iteration. E.g. if batchSize == 1, maxNumPaths == 2, maxPathLen== 3, [[[0, 1, 2], [0, 1, 10]]] - TensorPtr unpackedNextDraftTokens; // [maxBatchSize, maxNumPaths, maxPathLen], on gpu - //! Indices of draft tokens in the compressed `nextFlatTokens` for the next iteration. - //! Using example above, [[[0, 1, 2], [0, 1, 3]]] - TensorPtr unpackedNextDraftIndices; // [maxBatchSize, maxNumPaths, maxPathLen], on gpu - //! Probabilities of the next draft tokens. - TensorPtr nextDraftProbs; // [maxBatchSize, maxNumPaths, maxPathDraftLen, vocabSize], on gpu - //! Baseline for the position ids. - TensorPtr positionIdsBase; // [maxBatchSize], on gpu - //! Randomly sampled data (between 0.f and 1.f) - TensorPtr randomDataSample; // [maxBatchSize], on gpu - //! Randomly sampled data (between 0.f and 1.f) - TensorPtr randomDataValidation; // [maxBatchSize, maxNumPaths, maxDraftPathLen], on gpu - //! Sampling temperature. - TensorPtr temperatures; // [maxBatchSize], on gpu - //! Next generation lengths. - TensorPtr generationLengths; // [maxBatchSize], on gpu - //! Next generation lengths on host. - TensorPtr generationLengthsHost; // [maxBatchSize], on pinned - //! Maximum number of generated tokens for the next step across whole batch - TensorPtr maxGenLengthHost; // [1], on pinned -}; - -class EagleOutputs : public SpeculativeDecodingOutputs -{ -public: - explicit EagleOutputs(TensorPtr outputIds) - : SpeculativeDecodingOutputs{std::move(outputIds)} - { - } - - //! Unpacked draft tokens - TensorPtr unpackedNextDraftTokens; // [maxBatchSize, maxDecodingDraftTokens], on gpu - //! Draft paths for the next iteration. - TensorPtr nextDraftPaths; // [maxBatchSize, maxDecodingTokens, maxPathLen], on gpu - //! Randomly sampled data (between 0.f and 1.f) - TensorPtr randomDataSample; // [maxBatchSize], on gpu - //! Randomly sampled data (between 0.f and 1.f) - TensorPtr randomDataValidation; // [maxBatchSize], on gpu - //! Sampling temperature. - TensorPtr temperatures; // [maxBatchSize], on gpu - //! Next generation lengths. - TensorPtr generationLengths; // [maxBatchSize], on gpu - //! Next generation lengths. - TensorPtr generationLengthsHost; // [maxBatchSize], on pinned - //! Request types for ctx stage of the EagleNet0 (filled with 0s). - TensorPtr eagleNetCtxRequestTypesHost; // [maxBatchSize], on pinned - //! Context lengths of the context EagleNet0. - TensorPtr eagleNetCtxContextLengthsHost; // [maxBatchSize], on pinned - //! Past kv lengths of the context EagleNet0. - TensorPtr eagleNetCtxPastKeyValueLengthsHost; // [maxBatchSize], on pinned - //! Request types for ctx stage of the EagleNetX (filled with 1s). - TensorPtr eagleNetGenRequestTypesHost; // [maxBatchSize], on pinned - //! Context lengths of the generation EagleNetX. - TensorPtr eagleNetGenContextLengthsHost; // [maxBatchSize], on pinned - //! Past kv lengths of the generation EagleNetX. - TensorPtr eagleNetGenPastKeyValueLengthsHost; // [maxBatchSize], on pinned -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp b/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp deleted file mode 100644 index 41a6e7fb600e..000000000000 --- a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp +++ /dev/null @@ -1,312 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "dynamicDecodeLayer.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/decodingKernels.h" -#include "tensorrt_llm/layers/layerUtils.h" -#include "tensorrt_llm/layers/layersFactory.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" - -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -size_t DynamicDecodeLayer::getWorkspaceSize() const noexcept -{ - size_t maxWorkspaceSize = 0; - for (auto const& layer : mLayers) - { - maxWorkspaceSize = std::max(maxWorkspaceSize, layer->getWorkspaceSize()); - } - return maxWorkspaceSize; -} - -template -DynamicDecodeLayer::DynamicDecodeLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - initialize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::initialize() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mOutputIdsPtrHost = mBufferManager->pinnedPool(ITensor::makeShape({}), TRTDataType::value); - mParentIdsPtrHost = mBufferManager->pinnedPool(ITensor::makeShape({}), TRTDataType::value); - mOutputIdsPtrDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize())}), TRTDataType::value); - mParentIdsPtrDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize())}), TRTDataType::value); - - allocateBuffer(); - - mCyclicStep = 0; - mRuntimeMaxSeqLen = 0; - mConfiguredBeamWidth = -1; - - if (!mDecodingMode.isAuto()) - { - mConfiguredBeamWidth = mDecoderDomain.getBeamWidth(); - initializeLayers(); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mZeroParentIdsDevice - = mBufferManager->gpu(ITensor::makeShape({2 * mDecoderDomain.getBatchSize()}), TRTDataType::value); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::initializeLayers() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mLayers = createLayers(mDecodingMode, mDecoderDomain, mBufferManager); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::disableLookahead(DecoderDomain const& decoderDomain, SizeType32 batchSize, - TensorConstPtr batchSlots, std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mDecodingMode = executor::DecodingMode::TopKTopP(); - mDecoderDomain = std::move(decoderDomain); - initializeLayers(); - if (batchSize > 0) - { - setup(batchSize, 1, batchSlots, baseSetupParams, workspace); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - workspace->setDeviceBatchSlots( - batchSlots); // Copy the input batch slots to device for faster access in devie usage (kernels). - - TLLM_CHECK_WITH_INFO(setupParams->decodingParams, "decodingParams for setup is not set"); - if (setupParams->decodingParams->outputLogProbs) - { - // FIXME: monotonically growing - mOutputLogProbs = std::any_of(setupParams->decodingParams->outputLogProbs->begin(), - setupParams->decodingParams->outputLogProbs->end(), - [this](bool outputLogProbs) { return this->mOutputLogProbs | outputLogProbs; }); - } - - if (mConfiguredBeamWidth == -1) - { - // This code is left only for Python runtime - // In C++ runtime given maxBeamWidth should always be equal to the runtime beamWidth - TLLM_CHECK(mDecodingMode.isAuto()); - mConfiguredBeamWidth = beamWidth; - mDecodingMode - = mConfiguredBeamWidth == 1 ? executor::DecodingMode::TopKTopP() : executor::DecodingMode::BeamSearch(); - initializeLayers(); - auto const workspaceSize = getWorkspaceSize(); - workspace->resize(workspaceSize); - } - - TLLM_CHECK_WITH_INFO((mConfiguredBeamWidth == 1 && beamWidth == 1) - || (mConfiguredBeamWidth > 1 && beamWidth > 1 && beamWidth <= mConfiguredBeamWidth), - "Decoder is configured with beam width %d, but %d was given", mConfiguredBeamWidth, beamWidth); - TLLM_CHECK_WITH_INFO(mConfiguredBeamWidth <= mDecoderDomain.getBeamWidth(), - "Decoder is created with max beam width %d, but %d was given", mDecoderDomain.getBeamWidth(), - mConfiguredBeamWidth); - - for (auto& layer : mLayers) - { - layer->setup(batchSize, beamWidth, batchSlots, baseSetupParams, workspace); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(DynamicDecodeLayer_forwardAsync); - - auto params = std::dynamic_pointer_cast(baseInputs); - - TLLM_CHECK_WITH_INFO( - mDecodingMode.isExplicitDraftTokens() || mDecodingMode.isEagle() || params->logits || params->logitsVec, - "If not Explicit Draft Tokens or Eagle mode, either logits or logitsVec have to be specified."); - TLLM_CHECK_WITH_INFO( - baseOutputs->sequenceLength.has_value(), "sequenceLength tensor is required in DynamicDecoderLayer."); - - auto const localDecoderDomain = getLocalDecoderDomain(params, mDecoderDomain); - auto const maxSeqLen = baseOutputs->outputIds->getDimension<-1>(); - - TLLM_CHECK_WITH_INFO((mConfiguredBeamWidth == 1 && localDecoderDomain.getBeamWidth() == 1) - || (mConfiguredBeamWidth > 1 && localDecoderDomain.getBeamWidth() > 1 - && localDecoderDomain.getBeamWidth() <= mConfiguredBeamWidth), - "Decoder is configured with beam width %d, but %d was given", mConfiguredBeamWidth, - localDecoderDomain.getBeamWidth()); - - if (mOutputIdsPtrHost->getSize() == 0) - { - mOutputIdsPtrHost->reshape( - ITensor::makeShape({static_cast(maxSeqLen), static_cast(mDecoderDomain.getBatchSize())})); - mParentIdsPtrHost->reshape( - ITensor::makeShape({static_cast(maxSeqLen), static_cast(mDecoderDomain.getBatchSize())})); - mRuntimeMaxSeqLen = maxSeqLen; - } - - mCyclicStep = mCyclicStep % mRuntimeMaxSeqLen; - //! Copy the input batch slots to device for faster access in devie usage (kernels). - workspace->setDeviceBatchSlots(params->batchSlots); - - prepareIdsPtrs(baseOutputs, params->batchSlots, localDecoderDomain.getBatchSize(), - localDecoderDomain.getBeamWidth(), maxSeqLen); - - for (auto& layer : mLayers) - { - layer->forwardAsync(baseOutputs, baseInputs, workspace); - } - - // Copy nextIds and transpose logits when needed - prepareOutputData(baseOutputs, workspace->getDeviceBatchSlots(), localDecoderDomain.getBatchSize(), - mDecoderDomain.getBatchSize(), baseOutputs->beamWidth, maxSeqLen, mDecoderDomain.getMaxDecodingTokens(), - mOutputLogProbs, getStream()); - - mCyclicStep += 1; - - sync_check_cuda_error(getStream()); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::forwardSync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(DynamicDecodeLayer_forwardSync); - - for (auto& layer : mLayers) - { - layer->forwardSync(baseOutputs, baseInputs, workspace); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::prepareIdsPtrs(std::shared_ptr const& outputs, - BufferConstPtr batchSlots, SizeType32 batchSize, SizeType32 beamWidth, SizeType32 maxSeqLen) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TensorPtr outputIdsPtrHostSlice = ITensor::at(mOutputIdsPtrHost, {mCyclicStep}); - TensorPtr parentIdsPtrHostSlice = ITensor::at(mParentIdsPtrHost, {mCyclicStep}); - auto outputIdsPtrHost = runtime::bufferCast(*outputIdsPtrHostSlice); - auto parentIdsPtrHost = runtime::bufferCast(*parentIdsPtrHostSlice); - auto const* batchSlotsPtr = bufferCast(*batchSlots); - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - auto const batchSlot = batchSlotsPtr[bi]; - outputIdsPtrHost[batchSlot] = bufferCast(*outputs->outputIds) + batchSlot * beamWidth * maxSeqLen; - - if (beamWidth > 1) - { - parentIdsPtrHost[batchSlot] - = bufferCast(*outputs->parentIds.value()) + batchSlot * beamWidth * maxSeqLen; - } - else - { - auto mZeroParentIdsDevicePtr = bufferCast(*mZeroParentIdsDevice); - parentIdsPtrHost[batchSlot] = mZeroParentIdsDevicePtr + bi * beamWidth * maxSeqLen; - } - } - - mBufferManager->copy(*outputIdsPtrHostSlice, *mOutputIdsPtrDevice); - mBufferManager->copy(*parentIdsPtrHostSlice, *mParentIdsPtrDevice); - outputs->outputIdsPtr = ITensor::slice(mOutputIdsPtrDevice, 0, batchSize); - outputs->outputIdsPtrHost = ITensor::slice(outputIdsPtrHostSlice, 0, batchSize); - outputs->parentIdsPtr = ITensor::slice(mParentIdsPtrDevice, 0, batchSize); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayer::prepareOutputData(std::shared_ptr const& outputs, - BufferConstPtr batchSlots, SizeType32 batchSize, SizeType32 maxBatchSize, SizeType32 beamWidth, - SizeType32 maxSeqLen, SizeType32 maxTokensPerStep, bool outputLogProbs, cudaStream_t stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto outputIdsPtrDevice = bufferCast(*mOutputIdsPtrDevice); - auto const numNewTokens = bufferCastOrNull(outputs->numNewTokens); - auto newTokensPtr = bufferCast(*outputs->newTokens); - auto sequenceLengthsPtr = bufferCast(*outputs->sequenceLength.value()); - auto const* batchSlotsPtr = bufferCast(*batchSlots); - - invokeCopyNextStepIds(newTokensPtr, outputIdsPtrDevice, sequenceLengthsPtr, numNewTokens, batchSlotsPtr, batchSize, - maxBatchSize, beamWidth, maxSeqLen, maxTokensPerStep, stream); - - // Transpose output log probs from [maxSeqLen, batchSize, beamWidth] to [batchSize, beamWidth, maxSeqLen] - if (outputLogProbs && outputs->outputLogProbsTiled) - { - auto logProbsMaxSeqLen = outputs->outputLogProbsTiled.value()->getDimension<0>(); - - auto outputLogProbsPtr = bufferCast(*outputs->outputLogProbs.value()); - auto outputLogProbsTiledPtr = bufferCast(*outputs->outputLogProbsTiled.value()); - invokeTransposeLogProbs(outputLogProbsPtr, outputLogProbsTiledPtr, sequenceLengthsPtr, batchSlotsPtr, batchSize, - maxBatchSize, beamWidth, logProbsMaxSeqLen, stream); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class DynamicDecodeLayer; -template class DynamicDecodeLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h b/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h deleted file mode 100644 index c04fc345798f..000000000000 --- a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h +++ /dev/null @@ -1,88 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/penaltyLayer.h" - -namespace tensorrt_llm::layers -{ - -template -class DynamicDecodeLayer : public BaseLayer -{ - using Base = BaseLayer; - -public: - DynamicDecodeLayer(executor::DecodingMode const& mode, DecoderDomain const& decodingDomain, - std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, - runtime::ITensor::SharedConstPtr batchSlots, std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - void forwardSync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - - void disableLookahead(DecoderDomain const& decoderDomain, SizeType32 batchSize, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace); - -private: - void allocateBuffer(); - - void initialize(); - void initializeLayers(); - - void prepareIdsPtrs(std::shared_ptr const& outputs, BufferConstPtr batchSlots, - runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, runtime::SizeType32 maxSeqLen); - void prepareOutputData(std::shared_ptr const& outputs, BufferConstPtr batchSlots, - runtime::SizeType32 batchSize, runtime::SizeType32 maxBatchSize, runtime::SizeType32 beamWidth, - runtime::SizeType32 maxSeqLen, runtime::SizeType32 maxTokensPerStep, bool outputLogProbs, cudaStream_t stream); - -private: - using Base::mDecoderDomain; - - std::vector> mLayers; - - executor::DecodingMode mDecodingMode; - - TensorPtr mZeroParentIdsDevice; - TensorPtr mOutputIdsPtrHost; - TensorPtr mParentIdsPtrHost; - TensorPtr mOutputIdsPtrDevice; - TensorPtr mParentIdsPtrDevice; - - bool mHasDiffRuntimeArgs{false}; - - bool mOutputLogProbs{false}; - - runtime::SizeType32 mCyclicStep{0}; - runtime::SizeType32 mRuntimeMaxSeqLen{0}; - runtime::SizeType32 mConfiguredBeamWidth{-1}; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/eagleDecodingLayer.cpp b/cpp/tensorrt_llm/layers/eagleDecodingLayer.cpp deleted file mode 100644 index 4dce40985409..000000000000 --- a/cpp/tensorrt_llm/layers/eagleDecodingLayer.cpp +++ /dev/null @@ -1,313 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "eagleDecodingLayer.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/workspace.h" -#include "tensorrt_llm/kernels/speculativeDecoding/common.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" - -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -EagleDecodingLayer::EagleDecodingLayer( - DecoderDomain const& decoderDomain, std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, std::move(bufferManager)) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - allocateBuffer(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSizeShape = ITensor::makeShape({mDecoderDomain.getBatchSize()}); - mTemperature = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mTemperatureDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mCurandStatesDevice = mBufferManager->gpu( - ITensor::makeShape({mDecoderDomain.getBatchSize(), sizeof(curandState_t)}), TRTDataType::value); - - mEagleNetCtxRequestTypes = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mEagleNetCtxContextLengths = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mEagleNetCtxPastKeyValueLengths = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mEagleNetGenRequestTypes = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mEagleNetGenContextLengths = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mEagleNetGenPastKeyValueLengths = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - SizeType32 constexpr NUM_BUFFERS{1}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = mDecoderDomain.getBatchSize() * sizeof(SizeType32); - mWorkspaceSize = calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(EagleDecodingLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - workspace->initializeDeviceCurandStates( - setupParams->randomSeed, batchSize, workspace->getDeviceBatchSlots(), mCurandStatesDevice); - - // Setup penalties. - FillBuffers const fillBuffers{batchSize, mDecoderDomain.getBatchSize(), mBufferManager}; - - auto constexpr fltMax = std::numeric_limits::max(); - auto constexpr fltEpsilon = std::numeric_limits::epsilon(); - - // Allow temp = 0 as it will be overwritten in Eagle's typical acceptance codes. - fillBuffers(setupParams->temperature, DefaultDecodingParams::getTemperature(), mTemperature, mTemperatureDevice, - batchSlots, std::make_pair(-fltEpsilon, fltMax), "temperature penalty"); - - fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::fillContextBuffers(SizeType32 batchSize, BufferConstPtr batchSlots, - EagleSetupParams const& setupParams, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - FillContextEagleParams params; - params.outputRandDataSample = bufferCast(*setupParams.randomDataSample); - params.outputTemperatures = bufferCast(*setupParams.temperatures); - - params.inputTemperatures = bufferCastOrNull(mTemperatureDevice); - params.inputCurandState = reinterpret_cast(bufferCastOrNull(mCurandStatesDevice)); - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.batchSize = batchSize; - - params.checkParams(); - - invokeFillContextEagleData(params, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(EagleDecodingLayer_forwardSyncCPU); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - auto outputs = std::dynamic_pointer_cast(baseOutputs); - - // Convert batch slots and seq slots to have -1 for the ctx requests not in the last chunk. - augmentBatchSlots(*outputs, *inputs, workspace); - - // Slice output ids, pos ids, next draft tokens. - unpackData(*outputs, *inputs, workspace); - - // Convert masks to packed masks per request. - convertToPackedMask(*outputs, *inputs, workspace); - - // Pack accepted paths for KV cache rewind. - packAcceptedPaths(*outputs, *inputs, workspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::augmentBatchSlots(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.localBatchSize; - auto const engineBatchSize = inputs.nextDraftLens->getDimension<0>(); - - auto* workspaceBytePtr = reinterpret_cast(workspace->getRawWorkspaceDevicePtr()); - size_t offset{0}; - - auto* augmentedSeqSlots = reinterpret_cast( - nextWorkspacePtr(workspaceBytePtr, offset, engineBatchSize * sizeof(SizeType32))); - - auto const* chunkedContextNextTokens = bufferCast(*inputs.chunkedContextNextTokens); - auto const* lastDraftLens = bufferCast(*inputs.lastDraftLens); - - invokeAugmentBatchSlots(augmentedSeqSlots, chunkedContextNextTokens, lastDraftLens, - bufferCast(*inputs.seqSlots), engineBatchSize, batchSize, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::unpackData(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const engineBatchSize = inputs.nextDraftLens->getDimension<0>(); - auto const maxSeqLen = outputs.outputIds->getDimension<-1>(); - - auto* workspaceBytePtr = reinterpret_cast(workspace->getRawWorkspaceDevicePtr()); - size_t offset{0}; - - auto const* augmentedSeqSlots = reinterpret_cast( - nextWorkspacePtr(workspaceBytePtr, offset, engineBatchSize * sizeof(SizeType32))); - - UnpackEagleDataParams params; - params.batchSlots = augmentedSeqSlots; - params.inputCurandState = reinterpret_cast(bufferCastOrNull(mCurandStatesDevice)); - - params.inputTemperatures = bufferCast(*mTemperatureDevice); - params.inputNextDraftTokens = bufferCast(*inputs.nextDraftTokens); - params.inputNextDraftLens = bufferCast(*inputs.nextDraftLens); - params.inputNextDraftPaths = bufferCast(*inputs.nextDraftPaths); - params.inputLastDraftTokens = bufferCast(*inputs.lastDraftTokens); - params.inputLastDraftLens = bufferCast(*inputs.lastDraftLens); - params.inputAcceptedTokens = bufferCast(*inputs.acceptedTokens); - params.inputAcceptedLens = bufferCast(*inputs.acceptedLens); - - params.outputIds = bufferCast(*outputs.outputIds); - params.outputNumNewTokens = bufferCast(*outputs.numNewTokens.value()); - params.outputSequenceLengths = bufferCast(*outputs.sequenceLength.value()); - // FIXME outputUnpackedNextDraftTokens is the same as outputNextDraftTokens. - // outputUnpackedNextDraftTokens is used in eagleBuffers and outputNextDraftTokens is used in the runtime - params.outputUnpackedNextDraftTokens = bufferCast(*outputs.unpackedNextDraftTokens); - params.outputNextDraftTokens = bufferCast(*outputs.nextDraftTokens); - params.outputNextDraftLengths = bufferCast(*outputs.nextDraftLengths); - params.outputNextGenerationLength = bufferCast(*outputs.generationLengths); - params.outputNextDraftPaths = bufferCast(*outputs.nextDraftPaths); - params.outputPrevDraftLengths = bufferCast(*outputs.prevDraftLengths); - params.outputPositionIds = bufferCast(*outputs.nextDraftPosIds); - - params.outputRandDataSample = bufferCast(*outputs.randomDataSample); - params.outputRandDataVerification = bufferCast(*outputs.randomDataValidation); - - params.outputTemperatures = bufferCast(*outputs.temperatures); - - params.outputEagleNetCtxRequestTypes = bufferCast(*mEagleNetCtxRequestTypes); - params.outputEagleNetCtxContextLengths = bufferCast(*mEagleNetCtxContextLengths); - params.outputEagleNetCtxPastKeyValueLengths = bufferCast(*mEagleNetCtxPastKeyValueLengths); - params.outputEagleNetGenRequestTypes = bufferCast(*mEagleNetGenRequestTypes); - params.outputEagleNetGenContextLengths = bufferCast(*mEagleNetGenContextLengths); - params.outputEagleNetGenPastKeyValueLengths = bufferCast(*mEagleNetGenPastKeyValueLengths); - - params.batchSize = engineBatchSize; - params.maxDecodingTokens = mDecoderDomain.getSpeculativeDecodingModule()->getMaxDecodingTokens(); - params.maxPathLength = mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen(); - params.maxSeqLen = maxSeqLen; - - params.checkParams(); - - invokeUnpackEagleData(params, getStream()); - - mBufferManager->copy(*mEagleNetCtxRequestTypes, *outputs.eagleNetCtxRequestTypesHost); - mBufferManager->copy(*mEagleNetCtxContextLengths, *outputs.eagleNetCtxContextLengthsHost); - mBufferManager->copy(*mEagleNetCtxPastKeyValueLengths, *outputs.eagleNetCtxPastKeyValueLengthsHost); - mBufferManager->copy(*mEagleNetGenRequestTypes, *outputs.eagleNetGenRequestTypesHost); - mBufferManager->copy(*mEagleNetGenContextLengths, *outputs.eagleNetGenContextLengthsHost); - mBufferManager->copy(*mEagleNetGenPastKeyValueLengths, *outputs.eagleNetGenPastKeyValueLengthsHost); - - mBufferManager->copy(*outputs.generationLengths, *outputs.generationLengthsHost); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::convertToPackedMask(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const engineBatchSize = inputs.nextDraftLens->getDimension<0>(); - auto const maxDecodingTokens = mDecoderDomain.getSpeculativeDecodingModule()->getMaxDecodingTokens(); - auto const maxPathLen = mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen(); - - auto* workspaceBytePtr = reinterpret_cast(workspace->getRawWorkspaceDevicePtr()); - size_t offset{0}; - auto const* augmentedSeqSlots = reinterpret_cast( - nextWorkspacePtr(workspaceBytePtr, offset, engineBatchSize * sizeof(SizeType32))); - - auto const* batchSlots = augmentedSeqSlots; - auto* packedMasksDevice = bufferCast(*outputs.packedMasks); - auto const* nextDraftPaths = bufferCast(*outputs.nextDraftPaths); - - invokeGetPackedMaskFromPath( - packedMasksDevice, batchSlots, nextDraftPaths, engineBatchSize, maxDecodingTokens, maxPathLen, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodingLayer::packAcceptedPaths(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.localBatchSize; - auto const engineBatchSize = inputs.nextDraftLens->getDimension<0>(); - - auto* workspaceBytePtr = reinterpret_cast(workspace->getRawWorkspaceDevicePtr()); - size_t offset{0}; - auto const* augmentedSeqSlots = reinterpret_cast( - nextWorkspacePtr(workspaceBytePtr, offset, engineBatchSize * sizeof(SizeType32))); - - auto const* numNewTokens = bufferCast(*outputs.numNewTokens.value()); - auto* numNewTokensCumSum = bufferCast(*outputs.numNewTokensCumSum); - auto* pathsOffsets = bufferCast(*outputs.pathsOffsets); - auto const* batchSlots = augmentedSeqSlots; - auto const* bestPathIndicesSlotsPtr = bufferCast(*inputs.acceptedPathIds); - auto const* lastDraftPathsSlotsPtr = bufferCast(*inputs.lastDraftPaths); - - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for EagleDecodingLayer"); - TLLM_CHECK_WITH_INFO(numNewTokens != nullptr, "Accepted lengths must be provided for EagleDecodingLayer"); - TLLM_CHECK_WITH_INFO(numNewTokensCumSum != nullptr, "numNewTokensCumSum must be provided for EagleDecodingLayer"); - TLLM_CHECK_WITH_INFO(pathsOffsets != nullptr, "pathsOffsets must be provided for EagleDecodingLayer"); - invokePackAcceptedPaths(numNewTokensCumSum, pathsOffsets, numNewTokens, bestPathIndicesSlotsPtr, - lastDraftPathsSlotsPtr, batchSlots, batchSize, engineBatchSize, - mDecoderDomain.getSpeculativeDecodingModule()->getMaxNumPaths(), - mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen(), true, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t EagleDecodingLayer::getWorkspaceSize() const noexcept -{ - return mWorkspaceSize; -} - -template class EagleDecodingLayer; -template class EagleDecodingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/eagleDecodingLayer.h b/cpp/tensorrt_llm/layers/eagleDecodingLayer.h deleted file mode 100644 index 933b8f81b33a..000000000000 --- a/cpp/tensorrt_llm/layers/eagleDecodingLayer.h +++ /dev/null @@ -1,86 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Decoding layer for EAGLE speculative decoding technique. -template -class EagleDecodingLayer : public BaseLayer -{ -public: - using Base = BaseLayer; - using PathsVec = std::vector>>; - - EagleDecodingLayer(DecoderDomain const& decoderDomain, std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - void allocateBuffer(); - - void fillContextBuffers(runtime::SizeType32 batchSize, BufferConstPtr batchSlots, - EagleSetupParams const& setupParams, std::shared_ptr const& workspace); - - void augmentBatchSlots(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace); - - void convertToPackedMask(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace); - - void packAcceptedPaths(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace); - - void unpackData(EagleOutputs const& outputs, EagleInputs const& inputs, - std::shared_ptr const& workspace); - -private: - using Base::mDecoderDomain; - - size_t mWorkspaceSize{0}; - - TensorPtr mTemperature; - - TensorPtr mCurandStatesDevice; - TensorPtr mTemperatureDevice; - - TensorPtr mEagleNetCtxRequestTypes; - TensorPtr mEagleNetCtxContextLengths; - TensorPtr mEagleNetCtxPastKeyValueLengths; - TensorPtr mEagleNetGenRequestTypes; - TensorPtr mEagleNetGenContextLengths; - TensorPtr mEagleNetGenPastKeyValueLengths; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp deleted file mode 100644 index aedeb731574f..000000000000 --- a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp +++ /dev/null @@ -1,303 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "explicitDraftTokensLayer.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/penaltyTypes.h" -#include "tensorrt_llm/kernels/speculativeDecoding/common.h" -#include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -ExplicitDraftTokensLayer::ExplicitDraftTokensLayer( - DecoderDomain const& decoderDomain, std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - allocateBuffer(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExplicitDraftTokensLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mTemperature - = mBufferManager->pinnedPool(ITensor::makeShape({mDecoderDomain.getBatchSize()}), TRTDataType::value); - - mWorkspaceSize = invokeScanReduceGenerationLengths( - mDecoderDomain.getBatchSize(), nullptr, nullptr, 0, nullptr, nullptr, getStream()); - - mCurandStatesDevice = mBufferManager->gpu( - ITensor::makeShape({mDecoderDomain.getBatchSize(), sizeof(curandState_t)}), TRTDataType::value); - auto const batchSizeShape = ITensor::makeShape({mDecoderDomain.getBatchSize()}); - mGenerationLengthInclusiveSum = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mMaxGenerationLength = mBufferManager->gpu(ITensor::makeShape({1}), TRTDataType::value); - mTemperatureDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mBestPathIndicesSlots = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mLastDraftIndicesSlots = mBufferManager->gpu(ITensor::makeShape({mDecoderDomain.getBatchSize() - * mDecoderDomain.getSpeculativeDecodingModule()->getMaxNumPaths() - * mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen()}), - TRTDataType::value); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExplicitDraftTokensLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExplicitDraftTokensLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - workspace->initializeDeviceCurandStates( - setupParams->randomSeed, batchSize, workspace->getDeviceBatchSlots(), mCurandStatesDevice); - - // Setup penalties. - FillBuffers const fillBuffers{batchSize, mDecoderDomain.getBatchSize(), mBufferManager}; - - // Set decoder dtype to WAR the lack of bf16 support in decoder. - if (!mDecoderDtype) - { - mDecoderDtype = setupParams->dtype; - } - - fillBuffers(setupParams->temperature, DefaultDecodingParams::getTemperature(), mTemperature, mTemperatureDevice, - batchSlots, getLimitsPenalty(DecodingPenaltyType::Temperature), "temperature penalty"); - - // Dispatch context buffer fill - if (mDecoderDtype == tensorrt_llm::DataType::kFLOAT) - { - fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); - } - else if (mDecoderDtype == tensorrt_llm::DataType::kHALF) - { - fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); - } - else if (mDecoderDtype == tensorrt_llm::DataType::kBF16) - { - fillContextBuffers<__nv_bfloat16>(batchSize, batchSlots, *setupParams, workspace); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExplicitDraftTokensLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExplicitDraftTokensLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - auto outputs = std::dynamic_pointer_cast(baseOutputs); - - // DO NOT CHANGE THE ORDER. - - // Convert masks to packed masks per request. - convertPackedMask(*outputs, *inputs, workspace); - - // Slice output ids, pos ids, next draft tokens. - if (mDecoderDtype == tensorrt_llm::DataType::kFLOAT) - { - splitInputDataToBatchSlots(*outputs, *inputs, workspace); - } - else if (mDecoderDtype == tensorrt_llm::DataType::kHALF) - { - splitInputDataToBatchSlots(*outputs, *inputs, workspace); - } - else if (mDecoderDtype == tensorrt_llm::DataType::kBF16) - { - splitInputDataToBatchSlots<__nv_bfloat16>(*outputs, *inputs, workspace); - } - - // Pack accepted paths for KV cache rewind. - packAcceptedPaths(*outputs, *inputs, workspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t ExplicitDraftTokensLayer::getWorkspaceSize() const noexcept -{ - return mWorkspaceSize; -} - -template -template -void ExplicitDraftTokensLayer::fillContextBuffers(SizeType32 batchSize, BufferConstPtr batchSlots, - ExplicitDraftTokensSetupParams const& setupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - FillContextExplicitDraftTokensParams params; - params.randDataSample = bufferCast(*setupParams.randomDataSample); - params.outputTemperatures = bufferCast(*setupParams.temperatures); - params.inputTemperatures = bufferCastOrNull(mTemperatureDevice); - params.curandState = reinterpret_cast(bufferCastOrNull(mCurandStatesDevice)); - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.batchSize = batchSize; - - params.checkParams(); - - invokeFillContextBuffers(params, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -template -void ExplicitDraftTokensLayer::splitInputDataToBatchSlots(ExplicitDraftTokensOutputs const& outputs, - ExplicitDraftTokensInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.localBatchSize; - auto const maxSeqLen = outputs.outputIds->getDimension<-1>(); - - ExtractExplicitDraftTokensParams params; - - params.outputIds = bufferCast(*outputs.outputIds); - params.outputPositionIdsBase = bufferCast(*outputs.positionIdsBase); - params.outputPositionIds = bufferCast(*outputs.nextDraftPosIds); - params.outputNextDraftTokens = bufferCast(*outputs.nextDraftTokens); - params.unpackedNextDraftTokens = bufferCast(*outputs.unpackedNextDraftTokens); - params.unpackedNextDraftIndices = bufferCast(*outputs.unpackedNextDraftIndices); - params.acceptedLengths = bufferCast(*outputs.numNewTokens.value()); - params.nextDraftLengths = bufferCast(*outputs.nextDraftLengths); - params.prevDraftLengths = bufferCast(*outputs.prevDraftLengths); - params.sequenceLengths = bufferCast(*outputs.sequenceLength.value()); - params.randDataSample = bufferCast(*outputs.randomDataSample); - params.randDataVerification = bufferCast(*outputs.randomDataValidation); - params.outputDraftProbs = bufferCast(*outputs.nextDraftProbs); - params.outputTemperatures = bufferCast(*outputs.temperatures); - params.outputGenerationLengths = bufferCast(*outputs.generationLengths); - params.outputBestPathIndices = bufferCast(*mBestPathIndicesSlots); - params.outputLastDraftIndices = bufferCast(*mLastDraftIndicesSlots); - - params.batchSlots = bufferCast(*inputs.seqSlots); - params.nextDraftTokens = bufferCast(*inputs.nextDraftTokens); - params.lastDraftTokens = bufferCast(*inputs.lastDraftTokens); - params.inputUnpackedNextDraftIndices = bufferCast(*inputs.nextDraftIndices); - params.bestPathLengths = bufferCast(*inputs.bestPathLengths); - params.bestPathIndices = bufferCast(*inputs.bestPathIndices); - params.inputPositionIdsBase = bufferCast(*inputs.positionIdsBase); - params.packedPositionIds = bufferCast(*inputs.packedPosIds); - params.nextFlatTokens = bufferCast(*inputs.nextFlatTokens); - params.nextDraftProbs = bufferCast(*inputs.nextDraftProbs); - params.lastGenerationLengths = bufferCastOrNull(inputs.lastGenerationLengths); - params.generationLengthInclusiveSum = bufferCast(*mGenerationLengthInclusiveSum); - params.lastDraftIndices = bufferCast(*inputs.lastDraftIndices); - params.inputTemperatures = bufferCast(*mTemperatureDevice); - params.curandState = reinterpret_cast(bufferCastOrNull(mCurandStatesDevice)); - params.batchSize = batchSize; - params.numPaths = mDecoderDomain.getSpeculativeDecodingModule()->getMaxNumPaths(); - params.maxPathLength = mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen(); - params.maxSeqLen = maxSeqLen; - params.vocabSize = mDecoderDomain.getVocabSizePadded(); - params.numContextRequests = batchSize - inputs.lastDraftTokens->getDimension<0>(); - params.numGenerationRequests = inputs.lastDraftTokens->getDimension<0>(); - - params.checkParams(); - - // Copy max generation length - mBufferManager->copy(*inputs.maxGenLengthDevice, *outputs.maxGenLengthHost); - - invokeExtractExplicitDraftTokens(params, getStream()); - - invokeCopyProbs(params, getStream()); - - // Copy generation lengths - mBufferManager->copy(*outputs.generationLengths, *outputs.generationLengthsHost); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExplicitDraftTokensLayer::convertPackedMask(ExplicitDraftTokensOutputs const& outputs, - ExplicitDraftTokensInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto batchSlots = bufferCast(*inputs.seqSlots); - auto masksDevice = bufferCast(*inputs.masks); - auto generationLengths = bufferCast(*inputs.generationLengths); - auto packedMasksDevice = bufferCast(*outputs.packedMasks); - - auto const batchSize = inputs.localBatchSize; - - auto generationLengthInclusiveSumPtr = bufferCastOrNull(mGenerationLengthInclusiveSum); - auto workSpaceDevicePtr = workspace->getRawWorkspaceDevicePtr(); - auto maxGenerationLengthPtr = bufferCastOrNull(mMaxGenerationLength); - invokeScanReduceGenerationLengths(batchSize, generationLengths, workSpaceDevicePtr, mWorkspaceSize, - generationLengthInclusiveSumPtr, maxGenerationLengthPtr, getStream()); - - invokeConvertMaskToPackedMask(batchSize, generationLengthInclusiveSumPtr, maxGenerationLengthPtr, masksDevice, - batchSlots, mDecoderDomain.getSpeculativeDecodingModule()->getMaxDecodingDraftTokens(), - mDecoderDomain.getSpeculativeDecodingModule()->getMaxDecodingTokens(), packedMasksDevice, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExplicitDraftTokensLayer::packAcceptedPaths(ExplicitDraftTokensOutputs const& outputs, - ExplicitDraftTokensInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.localBatchSize; - - auto numNewTokens = bufferCast(*outputs.numNewTokens.value()); - auto numNewTokensCumSum = bufferCast(*outputs.numNewTokensCumSum); - auto pathsOffsets = bufferCast(*outputs.pathsOffsets); - auto batchSlots = workspace->getDeviceBatchSlotsPtr(); - auto bestPathIndicesSlotsPtr = bufferCastOrNull(mBestPathIndicesSlots); - auto lastDraftIndicesSlotsPtr = bufferCastOrNull(mLastDraftIndicesSlots); - - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for ExplicitDraftTokensLayer"); - TLLM_CHECK_WITH_INFO(numNewTokens != nullptr, "Accepted lengths must be provided for ExplicitDraftTokensLayer"); - TLLM_CHECK_WITH_INFO( - numNewTokensCumSum != nullptr, "numNewTokensCumSum must be provided for ExplicitDraftTokensLayer"); - TLLM_CHECK_WITH_INFO(pathsOffsets != nullptr, "pathsOffsets must be provided for ExplicitDraftTokensLayer"); - invokePackAcceptedPaths(numNewTokensCumSum, pathsOffsets, numNewTokens, bestPathIndicesSlotsPtr, - lastDraftIndicesSlotsPtr, batchSlots, batchSize, batchSize, - mDecoderDomain.getSpeculativeDecodingModule()->getMaxNumPaths(), - mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen(), false, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class ExplicitDraftTokensLayer; -template class ExplicitDraftTokensLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h deleted file mode 100644 index 17fca4513cf1..000000000000 --- a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h +++ /dev/null @@ -1,90 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Decoding layer for speculative decoding technique, when all tokens are generated, decoded and accepted in the -//! engine. -template -class ExplicitDraftTokensLayer : public BaseLayer -{ -public: - using Base = BaseLayer; - using PathsVec = std::vector>>; - - ExplicitDraftTokensLayer(DecoderDomain const& decoderDomain, std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - void allocateBuffer(); - - void convertPackedMask(ExplicitDraftTokensOutputs const& outputs, ExplicitDraftTokensInputs const& inputs, - std::shared_ptr const& workspace); - - void packAcceptedPaths(ExplicitDraftTokensOutputs const& outputs, ExplicitDraftTokensInputs const& inputs, - std::shared_ptr const& workspace); - - template - void fillContextBuffers(SizeType32 batchSize, BufferConstPtr batchSlots, - ExplicitDraftTokensSetupParams const& setupParams, - std::shared_ptr const& workspace); - - template - void splitInputDataToBatchSlots(ExplicitDraftTokensOutputs const& outputs, ExplicitDraftTokensInputs const& inputs, - std::shared_ptr const& workspace); - -private: - using Base::mDecoderDomain; - - SizeType32 mNumPaths; - SizeType32 mMaxPathLength; - - size_t mWorkspaceSize{0}; - - TensorPtr mCurandStatesDevice; - TensorPtr mGenerationLengthInclusiveSum; - TensorPtr mMaxGenerationLength; - TensorPtr mTemperatureDevice; - TensorPtr mBestPathIndicesSlots; - TensorPtr mLastDraftIndicesSlots; - - TensorPtr mTemperature; - - std::optional mDecoderDtype{std::nullopt}; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/externalDraftTokensLayer.cpp b/cpp/tensorrt_llm/layers/externalDraftTokensLayer.cpp deleted file mode 100644 index 568e78cf42ff..000000000000 --- a/cpp/tensorrt_llm/layers/externalDraftTokensLayer.cpp +++ /dev/null @@ -1,615 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "externalDraftTokensLayer.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/samplingTopPKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -#include - -namespace tksd = tensorrt_llm::kernels::speculative_decoding; - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -ExternalDraftTokensLayer::ExternalDraftTokensLayer(executor::DecodingMode const& mode, - DecoderDomain const& decoderDomain, std::shared_ptr bufferManager, bool isDeterministic, - bool isAirTopP) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) - , mIsDeterministic(isDeterministic) - , mIsAirTopP(isAirTopP) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO(!mDecodingMode.isBeamSearch(), "ExternalDraftTokensLayer does not support Beam search mode"); - - auto const deviceId = getDevice(); - TLLM_CUDA_CHECK(cudaGetDeviceProperties(&mDeviceProp, deviceId)); - - allocateBuffer(decoderDomain.getBatchSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::allocateBuffer(SizeType32 batchSize) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // top k workspace size - auto workspaceSize = getTopKWorkspaceSize(batchSize, 1, TOP_K_MAX, mDecoderDomain.getVocabSizePadded()); - mWorkspaceSize = std::max(workspaceSize, mWorkspaceSize); - // top p workspace size - workspaceSize = getTopPWorkspaceSize(batchSize, mDecoderDomain.getVocabSizePadded()); - mWorkspaceSize = std::max(workspaceSize, mWorkspaceSize); - - // multinomial (top p == 1) workspace size - workspaceSize = getAirTopPWorkspaceSize(batchSize, mDecoderDomain.getVocabSizePadded(), mIsDeterministic); - mWorkspaceSize = std::max(workspaceSize, mWorkspaceSize); - - // batchsize here is maxBatchSize - auto const batchSizeShape = ITensor::makeShape({batchSize}); - - mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({batchSize, sizeof(curandState_t)}), TRTDataType::value); - mBatchIsAccepted = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mRuntimeMultinomialDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - // host buffers. - mSkipTopKDecodeDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mSkipTopKDecodeHost = BufferManager::pinnedPool(batchSizeShape, TRTDataType::value); - mSkipTopPDecodeDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mSkipTopPDecodeHost = BufferManager::pinnedPool(batchSizeShape, TRTDataType::value); - auto skipTopPDecodeHostRange = BufferRange(*mSkipTopPDecodeHost); - std::fill(skipTopPDecodeHostRange.begin(), skipTopPDecodeHostRange.end(), true); - - mOutputIdsAfterSampling = mBufferManager->gpu( - ITensor::makeShape({batchSize, mDecoderDomain.getVocabSizePadded()}), TRTDataType::value); - mOutputIdsAfterSamplingPtrsHost = BufferManager::pinned(batchSizeShape, TRTDataType::value); - mOutputIdsAfterSamplingPtrsDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mTargetOutputIds = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - mRuntimeTopKDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mRuntimeTopKHost = BufferManager::cpu(batchSizeShape, TRTDataType::value); - - mRuntimeTopPDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - mReturnAllSelectedTokensPerSlotHost = BufferManager::pinned(batchSizeShape, TRTDataType::value); - mReturnAllSelectedTokensPerSlotDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - mMaskBuffer = mBufferManager->gpu( - ITensor::makeShape({batchSize, mDecoderDomain.getVocabSizePadded()}), TRTDataType::value); - - mSetupWorkspaceSize = std::max({mBatchIsAccepted->getSizeInBytes(), mRuntimeMultinomialDevice->getSizeInBytes(), - mOutputIdsAfterSampling->getSizeInBytes(), mTargetOutputIds->getSizeInBytes(), mMaskBuffer->getSizeInBytes()}); - - mTargetLogits = mBufferManager->gpu( - ITensor::makeShape({batchSize, mDecoderDomain.getVocabSizePadded()}), TRTDataType::value); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - workspace->initializeDeviceCurandStates( - setupParams->randomSeed, batchSize, workspace->getDeviceBatchSlots(), mCurandStatesDevice); - - auto& runtimeMultinomialDeviceTensor = const_cast(*mRuntimeMultinomialDevice); - tensorrt_llm::runtime::kernels::invokeFill(runtimeMultinomialDeviceTensor, 1.0f, mBufferManager->getStream()); - - // Prepare runtime top K - auto runtimeTopK = setupParams->runtimeTopK.value_or(std::vector{DefaultDecodingParams::getTopK()}); - auto runtimeTopP = setupParams->runtimeTopP.value_or(std::vector{DefaultDecodingParams::getTopP()}); - - auto const paramsSize = expandMatchElements(batchSize, runtimeTopK, runtimeTopP); - - TLLM_CHECK_WITH_INFO(paramsSize != 0, - fmtstr("ExternalDraftTokensLayer got parameter with unexpected size, want 1 or batchSize(%d), got" - "runtimeTopK.size() = %zu, runtimeTopP.size() = %zu", - batchSize, runtimeTopK.size(), runtimeTopP.size())); - - for (size_t i = 0; i < paramsSize; ++i) - { - auto& topK = runtimeTopK[i]; - auto& topP = runtimeTopP[i]; - clampTopK(topK); - clampTopP(topP); - regularizeTopKTopP(topK, topP); - } - - // Update parameters on both device and host, so we can - // - determine whether we can skip launch TopK / TopP kernel by examine mSkipTopKDecodeHost / mSkipTopPDecodeHost - // - select best kernel by examine mRuntimeTopKHost - // without consulting device memory, or we'll have to do an expensive synchronization. - SizeType32* topKsPtr = nullptr; - float* topPsPtr = nullptr; - - if (paramsSize > 1) - { - auto initWorkspaceSizes = getTopKInitWorkspaceSizes(batchSize); - calcAlignedPointers(workspace->getRawWorkspaceDevicePtr(), initWorkspaceSizes)(topKsPtr, topPsPtr); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, runtimeTopK, IBuffer::wrap(topKsPtr, initWorkspaceSizes[0] / sizeof(*topKsPtr))); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, runtimeTopP, IBuffer::wrap(topPsPtr, initWorkspaceSizes[1] / sizeof(*topPsPtr))); - } - auto const* batchSlotsDevicePtr = workspace->getDeviceBatchSlotsPtr(); - auto* skipTopKDecodeDevicePtr = bufferCastOrNull(mSkipTopKDecodeDevice); - auto* skipTopPDecodeDevicePtr = bufferCastOrNull(mSkipTopPDecodeDevice); - invokeSetupTopKTopPRuntimeArgs(batchSize, // - {topKsPtr, runtimeTopK.front(), bufferCast(*mRuntimeTopKDevice)}, // - {topPsPtr, runtimeTopP.front(), bufferCast(*mRuntimeTopPDevice)}, // - skipTopKDecodeDevicePtr, skipTopPDecodeDevicePtr, batchSlotsDevicePtr, true, getStream()); - - auto const* batchSlotsHostPtr = bufferCast(*batchSlots); - auto* skipDecodeTopKHostPtr = bufferCastOrNull(mSkipTopKDecodeHost); - auto* skipDecodeTopPHostPtr = bufferCastOrNull(mSkipTopPDecodeHost); - topKsPtr = paramsSize > 1 ? runtimeTopK.data() : nullptr; - invokeSetupTopKTopPRuntimeArgs(batchSize, // - {topKsPtr, runtimeTopK.front(), bufferCast(*mRuntimeTopKHost)}, {}, // - skipDecodeTopKHostPtr, skipDecodeTopPHostPtr, batchSlotsHostPtr, false); - - if (mIsAirTopP) - { - auto smCnt = mDeviceProp.multiProcessorCount; - if (smCnt <= 0) - { - auto const deviceId = getDevice(); - cudaDeviceProp prop{}; - TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, deviceId)); - smCnt = prop.multiProcessorCount; - } - mAirTopPBlockNum - = calcAirTopPBlockNum(batchSize, mDecoderDomain.getVocabSizePadded(), smCnt, mIsDeterministic); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::prepareInputs( - std::shared_ptr const& outputs, std::shared_ptr const& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // Fill the buffer for selected ids from sampling with zero. -1 will be set as a boundary if topP kernel is required - auto& outputIdsAfterSamplingTensor = const_cast(*mOutputIdsAfterSampling); - mBufferManager->setZero(outputIdsAfterSamplingTensor); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - if (inputs->step == 0) - { - // Prepare mReturnAllSelectedTokensPerSlot - auto numDraftTokensHost = BufferRange(*inputs->numDraftTokensHost); - auto returnAllSelectedTokensPerSlot = BufferRange(*mReturnAllSelectedTokensPerSlotHost); - std::transform(numDraftTokensHost.begin(), numDraftTokensHost.end(), returnAllSelectedTokensPerSlot.begin(), - [](auto numDraftTokens) { return numDraftTokens > 0; }); - mBufferManager->copy(*mReturnAllSelectedTokensPerSlotHost, *mReturnAllSelectedTokensPerSlotDevice); - - // Prepare mOutputIdsAfterSamplingPtrs - auto outputIdsAfterSamplingPtrsHost = BufferRange(*mOutputIdsAfterSamplingPtrsHost); - auto outputIdsPtrs = BufferRange(*outputs->outputIdsPtrHost); - - auto const maxBatchSize = mDecoderDomain.getBatchSize(); - for (auto batchSlot = 0; batchSlot < maxBatchSize; ++batchSlot) - { - auto outputIdsAfterSamplingSlice = ITensor::slice(mOutputIdsAfterSampling, batchSlot); - auto* outputIdsAfterSamplingPtr = bufferCast(*outputIdsAfterSamplingSlice); - - outputIdsAfterSamplingPtrsHost[batchSlot] - = returnAllSelectedTokensPerSlot[batchSlot] ? outputIdsAfterSamplingPtr : outputIdsPtrs[batchSlot]; - } - mBufferManager->copy(*mOutputIdsAfterSamplingPtrsHost, *mOutputIdsAfterSamplingPtrsDevice); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -}; - -template -void ExternalDraftTokensLayer::forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_forwardAsync); - - targetSoftmax(baseInputs, workspace); - - prepareInputs(outputs, baseInputs); - - // The logits from target engine should go through samplings first. - // gptDecoderBatched.cpp is calling dynamic decoder step by step, in this step, dynamic Decoder already forwarded - // PenaltyLayer, BanWordsLayer. For (TopK > 0) && (TopK == 0 && TopP == 0), we invoke TopK sampling kernel. The same - // logic is implemented in SamplingLayer.cpp - getAllTopKs(outputs, baseInputs, workspace); - - // Only for (TopK == 0 && TopP > 0), we invoke TopP sampling - getAllTopPs(outputs, baseInputs, workspace); - - // After all selected tokens are filled in mOutputIdsAfterSampling by topK, topP kernels, token acceptance logics - // starts. First we mask the logits of unselected token id to -inf as HF's TopK, TopP implementation. We compute the - // logit probs of draft and target and go through acceptance logics. - acceptDraftTokens(outputs, baseInputs, workspace); - - // If the token of the sequence is not accepted, a multinomial sampling is required for the bonus token. - // Multinomial sampling is achieved through TopP kernel with TopP = 1 and already weighted-sum target logits. - // The acceptance result of each batch is used as skipDecode in topP kernel. If is accepted, no sampling is needed - // (early exit). Forwarding for the next step is also set in this kernel. - multinomialSampling(outputs, baseInputs, workspace); - - // For the sequence with accepted tokens, we simply forward a step. - forwardAcceptedTokens(outputs, baseInputs, workspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t ExternalDraftTokensLayer::getWorkspaceSize() const noexcept -{ - return std::max(mWorkspaceSize, mSetupWorkspaceSize); -} - -template -void ExternalDraftTokensLayer::targetSoftmax(std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto const batchSize = inputs->logits.value()->getDimension<0>(); - - auto const* endIds = bufferCast(*inputs->endIds); - - FinishedState const* finishedInput = (inputs->finished) - ? reinterpret_cast(bufferCast(*inputs->finished.value())) - : nullptr; - - inputs->curandStates = reinterpret_cast(bufferCast(*mCurandStatesDevice)); - inputs->probsComputed = true; - - auto runtimeLogitsPtr = bufferCast(*workspace->getDeviceRuntimeLogits()); - auto logitsPtrsPtr = static_cast(nullptr); - auto biasPtr = static_cast(nullptr); - auto const* batchSlotsPtr = workspace->getDeviceBatchSlotsPtr(); - mBufferManager->copy(runtimeLogitsPtr, *mTargetLogits); - - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logits = runtimeLogitsPtr; - biasSoftmaxParams.logitsPtrs = logitsPtrsPtr; - biasSoftmaxParams.probs = runtimeLogitsPtr; - biasSoftmaxParams.bias = biasPtr; - biasSoftmaxParams.endIds = endIds; - biasSoftmaxParams.finished = finishedInput; - biasSoftmaxParams.batchSlots = batchSlotsPtr; - biasSoftmaxParams.batchSize = batchSize; - biasSoftmaxParams.maxBatchSize = mDecoderDomain.getBatchSize(); - biasSoftmaxParams.maxBeamWidth = 1; - biasSoftmaxParams.vocabSize = mDecoderDomain.getVocabSize(); - biasSoftmaxParams.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - biasSoftmaxParams.skipSoftMax = false; - biasSoftmaxParams.batchSlotsLogits = false; - biasSoftmaxParams.checkParams(); - - invokeAddBiasSoftMax(biasSoftmaxParams, getStream()); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::acceptDraftTokens(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_acceptDraftTokens); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto const draftLogitsShape = (*inputs->draftLogits).getShape(); - auto const maxBatchSize = mDecoderDomain.getBatchSize(); - auto const maxTokensPerStep = draftLogitsShape.d[1]; // 1 - auto const batchSize = static_cast(inputs->logits.value()->getDimension<0>()); - auto constexpr beamWidth = 1; - - FinishedState const* finishedInput = (inputs->finished) - ? reinterpret_cast(bufferCastOrNull(inputs->finished)) - : nullptr; - - FinishedState* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished)) - : nullptr; - - tksd::invokeMaskTargetLogits(batchSize, bufferCast(*mTargetLogits), workspace->getDeviceBatchSlotsPtr(), - beamWidth, mDecoderDomain.getVocabSizePadded(), finishedInput, maxBatchSize, - bufferCast(*mOutputIdsAfterSampling), bufferCastOrNull(mRuntimeTopKDevice), - bufferCast(*mMaskBuffer), getStream()); - - auto const* batchSlotsHost = bufferCast(*inputs->batchSlots); - auto const* useDraftLogitsHostPtr = bufferCastOrNull(inputs->useDraftLogitsHost); - auto const skipDraftLogits = allOfBatchSlots(batchSlotsHost, useDraftLogitsHostPtr, batchSize, false); - - if (!skipDraftLogits && inputs->step == 0) - { - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logits = bufferCast(*inputs->draftLogits); - biasSoftmaxParams.probs = bufferCast(*inputs->draftProbs); - biasSoftmaxParams.finished = finishedInput; - biasSoftmaxParams.batchSlots = workspace->getDeviceBatchSlotsPtr(); - biasSoftmaxParams.batchSize = batchSize; - biasSoftmaxParams.maxBatchSize = maxBatchSize; - biasSoftmaxParams.maxBeamWidth = beamWidth * maxTokensPerStep; - biasSoftmaxParams.vocabSize = mDecoderDomain.getVocabSize(); - biasSoftmaxParams.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - biasSoftmaxParams.skipSoftMax = false; - biasSoftmaxParams.batchSlotsLogits = true; - biasSoftmaxParams.checkParams(); - invokeAddBiasSoftMax(biasSoftmaxParams, getStream()); - } - - { - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logits = bufferCast(*mTargetLogits); - biasSoftmaxParams.probs = bufferCast(*inputs->targetProbs); - biasSoftmaxParams.finished = finishedInput; - biasSoftmaxParams.batchSlots = workspace->getDeviceBatchSlotsPtr(); - biasSoftmaxParams.batchSize = batchSize; - biasSoftmaxParams.maxBatchSize = maxBatchSize; - biasSoftmaxParams.maxBeamWidth = beamWidth; - biasSoftmaxParams.vocabSize = mDecoderDomain.getVocabSize(); - biasSoftmaxParams.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - biasSoftmaxParams.skipSoftMax = false; - biasSoftmaxParams.batchSlotsLogits = false; - biasSoftmaxParams.checkParams(); - invokeAddBiasSoftMax(biasSoftmaxParams, getStream()); - } - - sync_check_cuda_error(getStream()); - - tksd::invokeAcceptDraftTokens(batchSize, bufferCast(*inputs->draftProbs), bufferCast(*inputs->targetProbs), - bufferCast(*inputs->numDraftTokens), bufferCast(*inputs->useDraftLogits), - bufferCast(*inputs->draftTokenIds), finishedInput, finishedOutput, inputs->curandStates, - workspace->getDeviceBatchSlotsPtr(), maxTokensPerStep, beamWidth, mDecoderDomain.getVocabSizePadded(), - inputs->useRandomAcceptanceThreshold, inputs->constantThreshold, inputs->step, - bufferCast(*mBatchIsAccepted), bufferCast(*mTargetOutputIds), getStream()); - - sync_check_cuda_error(getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::multinomialSampling(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_multinomialSampling); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto const batchSize = inputs->logits.value()->getDimension<0>(); - auto probs = bufferCastOrNull(inputs->targetProbs); - auto* sequenceLength = bufferCastOrNull(outputs->sequenceLength); - auto const* endIds = bufferCastOrNull(inputs->endIds); - - FinishedState* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished)) - : nullptr; - - TopPSamplingKernelParams params{}; - params.probs = probs; - params.outputIdsPtrs = bufferCastOrNull(outputs->outputIdsPtr); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.topPs = bufferCastOrNull(mRuntimeMultinomialDevice); - params.sequenceLength = sequenceLength; - params.endIds = endIds; - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.finishedInput = nullptr; - params.finishedOutput = finishedOutput; - params.skipDecode = bufferCastOrNull(mBatchIsAccepted); - params.cumLogProbs = nullptr; - params.outputLogProbs = nullptr; - params.curandState = inputs->curandStates; - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - - if (!mIsAirTopP) - { - invokeBatchTopPSampling(params, getStream()); - } - else - { - params.blockNum = mAirTopPBlockNum; - params.isDeterministic = mIsDeterministic; - invokeBatchAirTopPSampling(params, getStream()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::getAllTopKs(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_getAllTopKs); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto logits = bufferCastOrNull(inputs->logits); - - auto const batchSize = static_cast(inputs->logits.value()->getDimension<0>()); - - auto const* batchSlotsHost = bufferCast(*inputs->batchSlots); - auto const* skipDecodeHostPtr = bufferCastOrNull(mSkipTopKDecodeHost); - auto const skip = allOfBatchSlots(batchSlotsHost, skipDecodeHostPtr, batchSize, true); - if (skip) - { - return; - } - - auto* sequenceLength = bufferCastOrNull(outputs->sequenceLength); - auto const* endIds = bufferCastOrNull(inputs->endIds); - - FinishedState const* finishedInput = (inputs->finished) - ? reinterpret_cast(bufferCastOrNull(inputs->finished)) - : nullptr; - FinishedState* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished)) - : nullptr; - - auto const* runtimeTopKHostPtr = bufferCast(*mRuntimeTopKHost); - - TopKSamplingKernelParams params{}; - params.logProbs = logits; - params.outputIdsPtrs = bufferCastOrNull(mOutputIdsAfterSamplingPtrsDevice); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.endIds = endIds; - params.sequenceLengths = sequenceLength; - params.maxTopP = 1.0F; - params.topPs = bufferCastOrNull(mRuntimeTopPDevice); - params.maxTopK = maxOfBatchSlots(batchSlotsHost, runtimeTopKHostPtr, batchSize); - params.topKs = bufferCastOrNull(mRuntimeTopKDevice); - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.finishedInput = finishedInput; - params.finishedOutput = finishedOutput; - params.skipDecode = bufferCastOrNull(mSkipTopKDecodeDevice); - params.curandState = inputs->curandStates; - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.maxTokensPerStep = 1; - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - params.returnAllSelectedTokens = true; - params.returnAllSelectedTokensPerSlot = bufferCastOrNull(mReturnAllSelectedTokensPerSlotDevice); - params.logitsHasProbs = inputs->probsComputed; - params.outputIdCurrentStep = bufferCastOrNull(mTargetOutputIds); - params.skipOutputIdCurrentStep = bufferCast(*inputs->useDraftLogits); - - invokeBatchTopKSampling(params, getStream()); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::getAllTopPs(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_getAllTopPs); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto logits = bufferCastOrNull(inputs->logits); - - auto const batchSize = static_cast(inputs->logits.value()->getDimension<0>()); - - auto const* batchSlotsHost = bufferCast(*inputs->batchSlots); - auto const* skipDecodeHostPtr = bufferCastOrNull(mSkipTopPDecodeHost); - auto const skip = allOfBatchSlots(batchSlotsHost, skipDecodeHostPtr, batchSize, true); - if (skip) - { - return; - } - - auto* sequenceLength = bufferCastOrNull(outputs->sequenceLength); - auto const* endIds = bufferCastOrNull(inputs->endIds); - - FinishedState const* finishedInput = (inputs->finished) - ? reinterpret_cast(bufferCastOrNull(inputs->finished)) - : nullptr; - FinishedState* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished)) - : nullptr; - - TopPSamplingKernelParams params{}; - params.probs = logits; - params.outputIdsPtrs = bufferCastOrNull(mOutputIdsAfterSamplingPtrsDevice); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.endIds = endIds; - params.sequenceLength = sequenceLength; - params.topPs = bufferCastOrNull(mRuntimeTopPDevice); - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.finishedInput = finishedInput; - params.finishedOutput = finishedOutput; - params.skipDecode = bufferCastOrNull(mSkipTopPDecodeDevice); - params.curandState = inputs->curandStates; - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - params.returnAllSelectedTokens = true; - params.returnAllSelectedTokensPerSlot = bufferCastOrNull(mReturnAllSelectedTokensPerSlotDevice); - params.outputIdCurrentStep = bufferCastOrNull(mTargetOutputIds); - params.skipOutputIdCurrentStep = bufferCast(*inputs->useDraftLogits); - - invokeBatchTopPSampling(params, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExternalDraftTokensLayer::forwardAcceptedTokens(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(ExternalDraftTokensLayer_forwardAcceptedTokens); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - auto const batchSize = inputs->logits.value()->getDimension<0>(); - - auto const draftLogitsShape = (*inputs->draftLogits).getShape(); - auto const maxTokensPerStep = draftLogitsShape.d[1]; // 1 - - FinishedState* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished)) - : nullptr; - - tksd::invokeForwardAcceptedTokens(batchSize, workspace->getDeviceBatchSlotsPtr(), - bufferCast(*mBatchIsAccepted), bufferCastOrNull(outputs->sequenceLength), - bufferCast(*inputs->draftTokenIds), bufferCastOrNull(outputs->outputIdsPtr), - inputs->step, maxTokensPerStep, bufferCastOrNull(inputs->endIds), finishedOutput, getStream()); - - sync_check_cuda_error(getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class ExternalDraftTokensLayer; -template class ExternalDraftTokensLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/externalDraftTokensLayer.h b/cpp/tensorrt_llm/layers/externalDraftTokensLayer.h deleted file mode 100644 index e3fd3149260d..000000000000 --- a/cpp/tensorrt_llm/layers/externalDraftTokensLayer.h +++ /dev/null @@ -1,110 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Top class for sampling layers. -//! It sets up and executes TopKSamplingLayer and TopPSamplingLayer samplings -template -class ExternalDraftTokensLayer : public BaseLayer -{ -public: - using Base = BaseLayer; - - ExternalDraftTokensLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager, bool isDeterministic = true, bool isAirTopP = true); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - using Base::mDecoderDomain; - - executor::DecodingMode mDecodingMode; - - size_t mWorkspaceSize{0}; - size_t mSetupWorkspaceSize{0}; - - TensorPtr mCurandStatesDevice; - TensorPtr mSkipTopKDecodeDevice; - TensorPtr mSkipTopKDecodeHost; - TensorPtr mSkipTopPDecodeDevice; - TensorPtr mSkipTopPDecodeHost; - - TensorPtr mBatchIsAccepted; - TensorPtr mRuntimeMultinomialDevice; - - TensorPtr mOutputIdsAfterSampling; - TensorPtr mOutputIdsAfterSamplingPtrsHost; - TensorPtr mOutputIdsAfterSamplingPtrsDevice; - TensorPtr mTargetOutputIds; - TensorPtr mRuntimeTopKDevice; - TensorPtr mRuntimeTopKHost; - TensorPtr mRuntimeTopPDevice; - TensorPtr mReturnAllSelectedTokensPerSlotHost; - TensorPtr mReturnAllSelectedTokensPerSlotDevice; - TensorPtr mMaskBuffer; - - TensorPtr mTargetLogits; - - // AirTopP - cudaDeviceProp mDeviceProp; - runtime::SizeType32 mAirTopPBlockNum{0}; - bool mIsDeterministic{true}; - bool mIsAirTopP{false}; - -private: - void allocateBuffer(runtime::SizeType32 batchSize); - void prepareInputs( - std::shared_ptr const& outputs, std::shared_ptr const& baseInputs); - void targetSoftmax(std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace); - void acceptDraftTokens(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace); - void multinomialSampling(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace); - void getAllTopKs(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace); - void getAllTopPs(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace); - void forwardAcceptedTokens(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace); -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/layerUtils.h b/cpp/tensorrt_llm/layers/layerUtils.h deleted file mode 100644 index 62a10cc11ca9..000000000000 --- a/cpp/tensorrt_llm/layers/layerUtils.h +++ /dev/null @@ -1,183 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include -#include -#include -#include -#include - -#include - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -namespace tensorrt_llm::layers -{ - -// Using a local lambda in beam search layers to fill buffers causes an internal compiler error on nvcc windows. -// As a workaround and to promote DRY, the fill logic is refactored into FillBuffers below. -struct FillBuffers -{ - using BufferPtr = runtime::IBuffer::SharedPtr; - using TensorConstPtr = runtime::ITensor::UniqueConstPtr; - using BufferConstPtr = runtime::IBuffer::SharedConstPtr; - - template - void operator()(std::optional> const& optParam, T const defaultValue, BufferPtr const& hostBuffer, - BufferPtr const& deviceBuffer, BufferConstPtr const& batchSlots, std::pair const& limits, - std::string const& name) const - { - // Specialize for `beamWidthArray` and `beamSearchSteps` - bool constexpr isVector = std::is_same_v>; - for (runtime::SizeType32 bi = 0; bi < batchSize; ++bi) - { - T value = defaultValue; - runtime::SizeType32 const batchSlot = runtime::bufferCast(*batchSlots)[bi]; - if (optParam) - { - if (optParam->size() == 1) - { - value = optParam->front(); - } - else - { - TLLM_CHECK_WITH_INFO( - optParam->size() == static_cast(batchSize), "Argument vector size mismatch."); - value = optParam->at(bi); - } - } - if constexpr (isVector) // Fill vector (beam width array) - { - size_t constexpr maxLength = tensorrt_llm::kernels::kMaxBeamWidthArrayLength; - auto hostBufferRange = runtime::BufferRange(*hostBuffer); - for (int i = 0; i < value.size(); ++i) - { - TLLM_CHECK_WITH_INFO( - limits.first < static_cast(value[i]) && static_cast(value[i]) <= limits.second, - "%s param (%f) is out of limits (%f, %f]", name.c_str(), static_cast(value[i]), - limits.first, limits.second); - hostBufferRange[batchSlot * maxLength + i] = value[i]; - } - for (int i = 0; i < maxLength - value.size(); ++i) - { - hostBufferRange[batchSlot * maxLength + value.size() + i] = value[value.size() - 1]; - } - } - else // Fill scalar - { - TLLM_CHECK_WITH_INFO( - limits.first < static_cast(value) && static_cast(value) <= limits.second, - "%s param (%f) is out of limits (%f, %f]", name.c_str(), static_cast(value), limits.first, - limits.second); - auto hostBufferRange = runtime::BufferRange(*hostBuffer); - hostBufferRange[batchSlot] = value; - } - } - - auto const hostSlice = runtime::IBuffer::slice(hostBuffer, 0, maxBatchSize); - auto deviceSlice = runtime::IBuffer::slice(deviceBuffer, 0, maxBatchSize); - mBufferManager->copy(*hostSlice, *deviceSlice); - } - - runtime::SizeType32 batchSize; - runtime::SizeType32 maxBatchSize; - std::shared_ptr mBufferManager; -}; - -template -bool allOfBatchSlots(runtime::SizeType32 const* batchSlotsHost, T const* data, runtime::SizeType32 batchSize, T value) -{ - return std::all_of( - batchSlotsHost, batchSlotsHost + batchSize, [&](runtime::SizeType32 b) { return data[b] == value; }); -} - -template -T maxOfBatchSlots(runtime::SizeType32 const* batchSlotsHost, T const* data, runtime::SizeType32 batchSize) -{ - return std::transform_reduce( - batchSlotsHost, batchSlotsHost + batchSize, std::numeric_limits::lowest(), - [](auto a, auto b) { return std::max(a, b); }, [&](auto i) { return data[i]; }); -} - -inline DecoderDomain getLocalDecoderDomain( - std::shared_ptr baseInputs, DecoderDomain const& globalDecoderDomain) -{ - auto inputs = std::dynamic_pointer_cast(baseInputs); - runtime::SizeType32 batchSize{baseInputs->localBatchSize}; - runtime::SizeType32 beamWidth{0}; - runtime::SizeType32 vocabSize{0}; - if (inputs->logits) - { - auto const& logitsShape = inputs->logits.value()->getShape(); - TLLM_CHECK(logitsShape.nbDims == 3 || logitsShape.nbDims == 4); - beamWidth = inputs->logits.value()->getDimension<-2>(); - vocabSize = inputs->logits.value()->getDimension<-1>(); - } - else if (inputs->logitsVec) - { - TLLM_CHECK(inputs->logitsVec->size()); - auto const& logitsShape = inputs->logitsVec.value()[0]->getShape(); - TLLM_CHECK(logitsShape.nbDims == 3 || logitsShape.nbDims == 4); - beamWidth = inputs->logitsVec.value()[0]->getDimension<-2>(); - vocabSize = inputs->logitsVec.value()[0]->getDimension<-1>(); - } - else if (inputs->batchSlots) - { - beamWidth = globalDecoderDomain.getBeamWidth(); - vocabSize = globalDecoderDomain.getVocabSize(); - } - else - { - TLLM_THROW("Can't get local Decoder domain"); - } - return {batchSize, beamWidth, vocabSize}; -} - -template -size_t expandMatchElements(size_t expandSize, std::vector&... vector) -{ - std::array vectorSizes{vector.size()...}; - - bool allSingle = true; - for (auto size : vectorSizes) - { - if (size == expandSize) - { - allSingle = false; - } - else if (size != 1) - { - return 0; - } - } - - if (allSingle) - { - return 1; - } - - (vector.resize(expandSize, vector.front()), ...); - return expandSize; -} - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/layersFactory.h b/cpp/tensorrt_llm/layers/layersFactory.h deleted file mode 100644 index 56c97b9ac87d..000000000000 --- a/cpp/tensorrt_llm/layers/layersFactory.h +++ /dev/null @@ -1,95 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/layers/banWordsLayer.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingLayer.h" -#include "tensorrt_llm/layers/penaltyLayer.h" -#include "tensorrt_llm/layers/stopCriteriaLayer.h" -#include -#include - -namespace tensorrt_llm::layers -{ -enum DecodingLayers_t -{ - PENALTY_LAYER, - BAN_WORDS_LAYER, - DECODING_LAYER, - STOP_CRITERIA_LAYER -}; - -static std::vector createDecodingLayerTypes(executor::DecodingMode const& mode) -{ - std::vector types = {}; - if (mode.isUsePenalty()) - { - types.push_back(DecodingLayers_t::PENALTY_LAYER); - } - if (mode.isUseBanWords()) - { - types.push_back(DecodingLayers_t::BAN_WORDS_LAYER); - } - types.push_back(DecodingLayers_t::DECODING_LAYER); - if (mode.isUseStopCriteria()) - { - types.push_back(DecodingLayers_t::STOP_CRITERIA_LAYER); - } - return types; -} - -template -static std::vector> createLayers(executor::DecodingMode const& mode, - DecoderDomain const& decodingDomain, std::shared_ptr const& bufferManager) -{ - std::vector> layers; - auto layerTypes = createDecodingLayerTypes(mode); - // Only when draft tokens and predicted and decoded by the engine, we can skip penalty layer. - if (!mode.isExplicitDraftTokens() && !mode.isEagle()) - { - TLLM_CHECK_WITH_INFO(layerTypes.size() && layerTypes[0] == DecodingLayers_t::PENALTY_LAYER, - "Penalty layer is required to be the first layer for any decoder configuration"); - } - for (auto&& type : layerTypes) - { - std::unique_ptr layer; - switch (type) - { - case DecodingLayers_t::PENALTY_LAYER: - layer = std::make_unique>(mode, decodingDomain, bufferManager); - break; - - case DecodingLayers_t::BAN_WORDS_LAYER: - layer = std::make_unique>(mode, decodingDomain, bufferManager); - break; - - case DecodingLayers_t::DECODING_LAYER: - layer = std::make_unique>(mode, decodingDomain, bufferManager); - break; - - case DecodingLayers_t::STOP_CRITERIA_LAYER: - layer = std::make_unique>(mode, decodingDomain, bufferManager); - break; - - default: TLLM_CHECK_WITH_INFO(false, "Unknown DecodingLayers_t"); break; - } - layers.push_back(std::move(layer)); - } - return layers; -} -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp b/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp deleted file mode 100644 index 76da89dfec0d..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp +++ /dev/null @@ -1,585 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/layers/lookaheadAlgorithm.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -#include -#include -#include - -namespace tensorrt_llm::layers -{ - -using namespace tensorrt_llm::runtime; - -LookaheadAlgorithm::LookaheadAlgorithm( - runtime::SizeType32 maxW, runtime::SizeType32 maxN, runtime::SizeType32 maxG, runtime::SizeType32 id) - : mPoolManager(maxG) - , mPrefillsMax(runtime::BufferManager::cpu( - runtime::ITensor::makeShape({(maxN <= 1 ? 0 : maxN - 2)}), tensorrt_llm::DataType::kINT32)) - , mPastTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW * (maxN - 1)}), tensorrt_llm::DataType::kINT32)) - , mKeyTokensMax(runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW}), tensorrt_llm::DataType::kINT32)) - , mGoldenTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxN * 2 - 1}), tensorrt_llm::DataType::kINT32)) - , mGuessTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxG * (maxN - 1)}), tensorrt_llm::DataType::kINT32)) - , mMaxW(maxW) - , mMaxN(maxN) - , mMaxG(maxG) - , mFilling(0) -{ - runtime::SizeType32 maxGeneratedLen, maxDraftLen; - std::tie(maxGeneratedLen, std::ignore, maxDraftLen, std::ignore) - = executor::LookaheadDecodingConfig(maxW, maxN, maxG).calculateSpeculativeResource(); - mAttentionMask = runtime::BufferManager::cpu( - runtime::ITensor::makeShape({maxDraftLen, maxDraftLen}), tensorrt_llm::DataType::kBOOL); - mDraftTokensMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), tensorrt_llm::DataType::kINT32); - mSampledTokensMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxGeneratedLen}), tensorrt_llm::DataType::kINT32); - mEncodeMapMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), tensorrt_llm::DataType::kINT32); -} - -void LookaheadAlgorithm::setup(TensorConstPtr const& prompt, SizeType32 w, SizeType32 n, SizeType32 g, uint64_t seed) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK_WITH_INFO(w <= mMaxW, "lookahead requires setup w (%d) <= max_w (%d)", w, mMaxW); - TLLM_CHECK_WITH_INFO(n <= mMaxN, "lookahead requires setup n (%d) <= max_n (%d)", n, mMaxN); - TLLM_CHECK_WITH_INFO(g <= mMaxG, "lookahead requires setup g (%d) <= max_g (%d)", g, mMaxG); - mW = w; - mN = n; - mG = g; - std::tie(std::ignore, std::ignore, mRuntimeMaxDraftLen, mRuntimeMaxDraftPathLen) - = executor::LookaheadDecodingConfig(mW, mN, mG).calculateSpeculativeResource(); - - mPoolManager.setup(mG); - mPoolManager.accept(prompt, mN); - mGoldenTokens = ITensor::slice(mGoldenTokensMax, 0, mN * 2 - 1); - mPrefills = ITensor::slice(mPrefillsMax, 0, mN <= 1 ? 0 : mN - 2); - mKeyTokens = ITensor::slice(mKeyTokensMax, 0, mW); - mPastTokens = ITensor::slice(mPastTokensMax, 0, mW * (mN - 1)); - mPastTokens->reshape(ITensor::makeShape({mW, mN - 1})); - - BufferRange promptRange(*prompt); - BufferRange prefillRange(*mPrefills); - BufferRange pastRange(*mPastTokens); - BufferRange goldRange(*mGoldenTokens); - - srand(seed); - - auto randToken = [&promptRange](auto& item) { item = promptRange[rand() % promptRange.size()]; }; - std::for_each(prefillRange.begin(), prefillRange.end(), randToken); - std::for_each(pastRange.begin(), pastRange.end(), [](auto& a) { a = -1; }); - for (SizeType32 i = 0; i < mW; i++) - { - if (mN - 1 > 0) - { - randToken(pastRange[i * (mN - 1)]); - } - } - std::copy(std::prev(promptRange.end(), mN - 1), promptRange.end(), goldRange.begin()); - mGuessTokens = ITensor::slice(mGuessTokensMax, 0, 0); - mFilling = (mN - 1) > 0 ? 1 : 0; - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadAlgorithm::accept(TensorConstPtr const& generatedTokens) -{ - TLLM_CHECK(ITensor::volume(generatedTokens->getShape()) <= mN); - BufferRange generatedRange(*generatedTokens); - BufferRange goldRange(*mGoldenTokens); - auto genLen = generatedTokens->getShape().d[0]; - TLLM_CHECK(genLen <= mN); - std::copy(generatedRange.begin(), generatedRange.end(), goldRange.begin() + mN - 1); - TensorPtr newGold = ITensor::slice(mGoldenTokens, 0, mN - 1 + genLen); - mPoolManager.accept(newGold, mN); - std::copy(goldRange.begin() + genLen, goldRange.begin() + genLen + mN - 1, goldRange.begin()); -} - -//! lookahead has two phase, prefill the past tokens matrix and maintain past tokens matrix. -runtime::SizeType32 LookaheadAlgorithm::lookahead( - TensorPtr const& draftTokens, TensorPtr const& positionIds, runtime::SizeType32 startPosId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 prefill = mN - 2 - mFilling; - SizeType32 len = prefill + mFilling * mW; - TLLM_CHECK(len <= ITensor::volume(draftTokens->getShape())); - TLLM_CHECK(len <= ITensor::volume(positionIds->getShape())); - BufferRange prefillRange(*mPrefills); - BufferRange pastRange(*mPastTokens); - BufferRange draftRange(*draftTokens); - PRINT_TOKENS(mPrefills); - - if (mFilling < mN - 1) - { // prefilling - std::copy(prefillRange.begin() + mFilling, prefillRange.end(), draftRange.begin()); - for (SizeType32 i = 0; i < mW; i++) - { - auto start = pastRange.begin() + i * (mN - 1); - auto end = pastRange.begin() + i * (mN - 1) + mFilling; - std::copy(start, end, draftRange.begin() + prefill + i * mFilling); - } - } - else - { // shift up - std::copy(pastRange.begin() + 1, pastRange.begin() + mFilling * mW, draftRange.begin()); - } - - BufferRange positionIdsRange(*positionIds); - SizeType32 idx = 0, wj = 0; - auto fillPosition = [&positionIdsRange, &idx](SizeType32 start, SizeType32 len) - { - for (SizeType32 i = start; i < start + len; i++) - { - positionIdsRange[idx++] = i; - } - }; - if (prefill >= 0) - { - fillPosition(startPosId, prefill); - for (wj = 0; wj < mW; wj++) - { - fillPosition(startPosId + prefill + wj, mFilling); - } - } - else - { - fillPosition(startPosId, mFilling - 1); - for (wj = 1; wj < mW; wj++) - { - fillPosition(startPosId - 1 + wj, mFilling); - } - } - PRINT_VALUES(positionIds); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return len; -} - -runtime::SizeType32 LookaheadAlgorithm::guess(TensorPtr const& guessTokens, TensorPtr const& guessIds, - runtime::SizeType32 startPosId, runtime::TokenIdType lastToken) -{ - auto guesses = mPoolManager.guess(lastToken, mW); - - SizeType32 len = 0; - std::for_each(guesses.begin(), guesses.end(), [&len](auto& a) { len += ITensor::volume(a->getShape()); }); - TLLM_CHECK(len <= ITensor::volume(guessTokens->getShape())); - TLLM_CHECK(len <= ITensor::volume(guessIds->getShape())); - BufferRange guessTokensRange(*guessTokens); - BufferRange guessIdsRange(*guessIds); - - SizeType32 cur = 0; - for (auto guess : guesses) - { - BufferRange guessRange(*guess); - std::copy(guessRange.begin(), guessRange.end(), guessTokensRange.begin() + cur); - SizeType32 tmp = startPosId; - std::for_each( - guessIdsRange.begin() + cur, guessIdsRange.begin() + cur + mN - 1, [&tmp](auto& v) { v = tmp++; }); - cur += ITensor::volume(guess->getShape()); - } - - return len; -} - -void LookaheadAlgorithm::posIdsToMask(TensorPtr const& mask, TensorConstPtr const& posIds) -{ - auto len = ITensor::volume(posIds->getShape()); - TLLM_CHECK(mask->getDimension<0>() >= len); - TLLM_CHECK(mask->getDimension<1>() >= len); - auto posIdsRange = BufferRange(*posIds); - auto maskLocation = BufferLocation(*mask); - - for (auto& item : maskLocation) - { - item = false; - } - - if (len > 0) - { - std::vector> stack; - for (auto i = 0; i < len; i++) - { - auto cur = posIdsRange[i]; - while (stack.size() > 0 && cur <= stack.back().second) - { - stack.pop_back(); - } - TLLM_CHECK(stack.size() > 0 ? cur == stack.back().second + 1 : true); - stack.push_back(std::make_pair(i, cur)); - for (auto prev : stack) - { - maskLocation.at(i, prev.first) = true; - } - } - } -} - -struct TreeValue; -using TreeMap = std::unordered_map; - -struct TreeValue -{ - TreeValue() - : nexts(std::make_shared()) - { - } - - using Nexts = std::shared_ptr; - Nexts nexts{nullptr}; - std::list sources; -}; - -using TreeNode = TreeMap::value_type; - -template -void treeDFS(TreeNode& node, BF const& visitBefore, AF const& visitAfter) -{ - visitBefore(node); - for (auto& next : *(node.second.nexts)) - { - treeDFS(next, visitBefore, visitAfter); - } - visitAfter(node); -} - -SizeType32 LookaheadAlgorithm::treeEncode( - TensorPtr const& tokens, TensorPtr const& posIds, TensorPtr const& mask, TensorPtr const& encodeMap) -{ - TLLM_CHECK(ITensor::volume(tokens->getShape()) == ITensor::volume(posIds->getShape())); - auto len = ITensor::volume(tokens->getShape()); - - BufferRange tokensRange(*tokens); - BufferRange posIdsRange(*posIds); - BufferLocation maskLocation(*mask); - BufferRange mapRange(*encodeMap); - - auto branches = std::make_shared(); - - for (auto i = 0; i < len; i++) - { - auto nexts = branches; - for (auto j = 0; j <= i; j++) - { - if (maskLocation.at(i, j)) - { - auto tok = tokensRange[j]; - auto found = nexts->find(tok); - if (found != nexts->end()) - { - found->second.sources.push_back(j); - nexts = found->second.nexts; - } - else - { - auto [inserted, ok] = nexts->insert({tok, TreeValue()}); - inserted->second.sources.push_back(j); - nexts = inserted->second.nexts; - } - } - } - } - - for (auto& item : maskLocation) - { - item = 0; - } - std::vector> stack; - SizeType32 offset = 0; - SizeType32 posId = posIdsRange.size() ? posIdsRange[0] : 0; - - auto visitBefore - = [&stack, &maskLocation, &tokensRange, &posIdsRange, &posId, &offset, &mapRange](TreeNode const& node) - { - stack.push_back(std::make_pair(offset, node.first)); - for (auto const& source : node.second.sources) - { - mapRange[source] = offset; - } - for (auto const& prev : stack) - { - maskLocation.at(offset, prev.first) = true; - } - tokensRange[offset] = node.first; - posIdsRange[offset] = posId; - offset++; - posId++; - }; - auto visitAfter = [&stack, &posId](TreeNode const& node) - { - stack.pop_back(); - posId--; - }; - - for (auto& next : *branches) - { - treeDFS(next, visitBefore, visitAfter); - } - - for (SizeType32 i = offset; i < len; i++) - { - tokensRange[i] = 0; - posIdsRange[i] = 0; - } - for (SizeType32 i = 0; i < len; i++) - { - for (SizeType32 j = i < offset ? offset : 0; j < len; j++) - { - maskLocation.at(i, j) = false; - } - } - - return offset; -} - -void LookaheadAlgorithm::prepare(TensorPtr const& draftTokens, TensorPtr const& positionIds, - TensorPtr const& draftLengthPtr, TensorPtr const& attentionMask, SizeType32 attentionMaskOffset, - TensorConstPtr const& lastPositionIdPtr, TensorConstPtr const& lastTokenPtr) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (mRuntimeMaxDraftLen == 0) - { - mDraftTokens = ITensor::slice(mDraftTokensMax, 0, 0); - mEncodeMap = ITensor::slice(mEncodeMapMax, 0, 0); - (BufferRange(*draftLengthPtr))[0] = 0; - return; - } - - auto lastToken = BufferRange(*lastTokenPtr)[0]; - auto offset = BufferRange(*lastPositionIdPtr)[0]; - - SizeType32 inputLen = ITensor::volume(draftTokens->getShape()); - TLLM_CHECK(inputLen >= mRuntimeMaxDraftLen); - - BufferRange draftRange(*draftTokens); - BufferRange positionRange(*positionIds); - - SizeType32 filledLen = 0; - - filledLen += lookahead(ITensor::slice(draftTokens, filledLen, mRuntimeMaxDraftLen - filledLen), - ITensor::slice(positionIds, filledLen, mRuntimeMaxDraftLen - filledLen), offset); - - auto guessStart = filledLen; - filledLen += guess(ITensor::slice(draftTokens, filledLen, mRuntimeMaxDraftLen - filledLen), - ITensor::slice(positionIds, filledLen, mRuntimeMaxDraftLen - filledLen), offset, lastToken); - auto guessEnd = filledLen; - - std::copy(draftRange.begin() + guessStart, draftRange.begin() + guessEnd, - BufferRange(*mGuessTokensMax).begin()); - mGuessTokens = ITensor::slice(mGuessTokensMax, 0, guessEnd - guessStart); - - posIdsToMask(mAttentionMask, ITensor::slice(positionIds, 0, filledLen)); - - auto draftLen = treeEncode(ITensor::slice(draftTokens, 0, filledLen), ITensor::slice(positionIds, 0, filledLen), - mAttentionMask, mEncodeMapMax); - - for (SizeType32 i = 0; i < draftLen; i++) - { - BufferRange srcRange(*ITensor::at(mAttentionMask, {i})); - BufferRange dstRange(*ITensor::slice(attentionMask, {i + attentionMaskOffset, attentionMaskOffset})); - std::copy(srcRange.begin(), srcRange.end(), dstRange.begin()); - } - - std::copy(draftRange.begin(), draftRange.begin() + draftLen, BufferRange(*mDraftTokensMax).begin()); - mDraftTokens = ITensor::slice(mDraftTokensMax, 0, draftLen); - (BufferRange(*draftLengthPtr))[0] = draftLen; - mEncodeMap = ITensor::slice(mEncodeMapMax, 0, filledLen); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadAlgorithm::verify(TensorPtr const& accepted, TensorPtr const& acceptedOffsets, - TensorPtr const& acceptedLength, TokenIdType newLastToken, TensorConstPtr const& goldenTokens, - TensorConstPtr const& endToken) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(ITensor::volume(goldenTokens->getShape()) == ITensor::volume(mDraftTokens->getShape())); - BufferRange goldRange(*goldenTokens); - BufferRange draftRange(*mDraftTokens); - BufferLocation maskLocation(*mAttentionMask); - auto draftSize = ITensor::volume(mDraftTokens->getShape()); - auto end = *BufferRange(*endToken).begin(); - - SizeType32 maxHit = 0, hitIdx = 0; - for (SizeType32 i = 0; i < draftSize; i++) - { - SizeType32 hit = 0; - TokenIdType cur = newLastToken; - for (SizeType32 j = 0; j < draftSize; j++) - { - if (maskLocation.at(i, j)) - { - if (draftRange[j] == cur && draftRange[j] != end) - { - hit++; - cur = goldRange[j]; - } - else - { - break; - } - } - } - if (hit > maxHit) - { - maxHit = hit; - hitIdx = i; - } - } - - maxHit = maxHit > mRuntimeMaxDraftPathLen ? mRuntimeMaxDraftPathLen : maxHit; - - SizeType32 acceptedIdx = 0; - BufferRange acceptedRange(*accepted); - BufferRange acceptedOffsetsRange(*acceptedOffsets); - acceptedRange[acceptedIdx] = newLastToken; - for (SizeType32 j = 0; j < draftSize; j++) - { - if (maskLocation.at(hitIdx, j) && acceptedIdx < maxHit) - { - acceptedOffsetsRange[acceptedIdx++] = j; - acceptedRange[acceptedIdx] = goldRange[j]; - } - } - - *BufferRange(*acceptedLength).begin() = maxHit + 1; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! lookahead Jacobi matrix has prefilling phase and maintenance phase. -//! W=5, N=5. -//! *prefilling phase* -//! mFilling = 1->2, Tokens initialized from prompt. To fill the second line. -//! 0>1 2 3 * -//! 4 * -//! 5 * -//! 6 * -//! 7 * -//! mFilling = 2->3. -//! 0 1>2 3 4 * -//! 4 5 * -//! 5 6 * -//! 6 7 * -//! 7 8 * -//! mFilling = 3->4. -//! 0 1 2>3 4 5 * -//! 4 5 6 * -//! 5 6 7 * -//! 6 7 9 * -//! 7 8 a * -//! *maintenance phase* -//! mFilling = 4->4. shift up and generate five n-grams. -//! 0 1 2 3>4 5 6 * -//! 4 5 6 7 * -//! 5 6 7 8 * -//! 6 7 8 9 * -//! 7 8 9 a * -//! mFilling = 4. -//! 0 1 2 3 4>5 6 7 * -//! 5 6 7 8 * -//! 6 7 8 9 * -//! 7 8 9 a * -//! 8 9 a b * -//! mFilling = 4. -//! 0 1 2 3 4 5>6 7 8 * -//! 6 7 8 9 * -//! 7 8 9 a * -//! 8 9 a b * -//! 9 a b c * -void LookaheadAlgorithm::update(TensorPtr const& acceptedTokens, TensorPtr const& acceptedOffsets, - TensorPtr const& acceptedLength, TensorConstPtr const& sampledTokens, TensorConstPtr const& endToken) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(ITensor::volume(acceptedTokens->getShape()) >= mN); - BufferRange zippedTokensRange(*sampledTokens); - BufferRange sampledRange(*mSampledTokensMax); - - BufferRange mapRange(*mEncodeMap); - BufferRange unzipRange(*mSampledTokensMax); - mSampledTokens = ITensor::slice(mSampledTokensMax, 0, mEncodeMap->getShape().d[0] + 1); - - unzipRange[0] = zippedTokensRange[0]; - for (size_t i = 0; i < mapRange.size(); i++) - { - unzipRange[i + 1] = zippedTokensRange[mapRange[i] + 1]; - } - - BufferRange keyRange(*mKeyTokens); - BufferRange pastRange(*mPastTokens); - - auto newLastToken = sampledRange[0]; - SizeType32 prefill = mN - 2 - mFilling; - for (SizeType32 i = 0; i < mW; i++) - { - keyRange[i] = sampledRange[prefill + i * mFilling + mFilling]; - } - - if (mFilling < mN - 1) - { - for (SizeType32 i = 0; i < mW; i++) - { - pastRange[i * (mN - 1) + mFilling] = keyRange[i]; - } - } - else if (mN > 1) - { - for (SizeType32 i = 0; i < mW; i++) - { - auto begin = pastRange.begin() + i * (mN - 1); - auto end = pastRange.begin() + i * (mN - 1) + mN - 1; - auto key = *begin; - std::copy(begin + 1, end, begin); - *(std::prev(end, 1)) = keyRange[i]; - keyRange[i] = key; - } - keyRange[0] = newLastToken; - mPoolManager.update(mKeyTokens, mPastTokens); - } - - auto guessSize = ITensor::volume(mGuessTokens->getShape()); - auto outputSize = ITensor::volume(mSampledTokens->getShape()); - auto lookSize = 1 + (mN > 1 ? mN - 2 : 0) - mFilling + mFilling * mW; - TLLM_CHECK(guessSize + lookSize == outputSize); - - TensorConstPtr goldenTokens = ITensor::slice(mSampledTokens, lookSize, guessSize); - - verify(acceptedTokens, acceptedOffsets, acceptedLength, newLastToken, ITensor::slice(sampledTokens, 1), endToken); - - accept(ITensor::slice(acceptedTokens, 0, *BufferRange(*acceptedLength).begin())); - - if (mFilling < mN - 1) - { - mFilling++; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.h b/cpp/tensorrt_llm/layers/lookaheadAlgorithm.h deleted file mode 100644 index 02f271788c25..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.h +++ /dev/null @@ -1,135 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "lookaheadPoolManager.h" -#include "tensorrt_llm/runtime/common.h" - -namespace tensorrt_llm::layers -{ - -//! @brief An CPU implementation of Lookahead with ITensor. -class LookaheadAlgorithm -{ -public: - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorConstPtr = runtime::ITensor::SharedConstPtr; - - //! @brief Currently the resource management is to be aligned with batch manager. - //! @param w, n, g is the Jacobi window, n-gram level and guess set size respectively. - LookaheadAlgorithm( - runtime::SizeType32 maxW, runtime::SizeType32 maxN, runtime::SizeType32 maxG, runtime::SizeType32 id = 0); - - //! @brief setup per request, fill internal states from @param prompt. - void setup(TensorConstPtr const& prompt, runtime::SizeType32 w, runtime::SizeType32 n, runtime::SizeType32 g, - uint64_t seed); - - //! @brief accept the new generated tokens. - //! LookaheadDecodingLayer need call once for the first token in generation phase. - void accept(TensorConstPtr const& generatedTokens); - - //! @brief combine lookahead and guess to prepare the tensors. - //! input @param lastPositionIdPtr is position id of the last golden token, in a TensorPtr. - //! input @param lastTokenPtr the last golden token for searching in the pool, in a TensorPtr. - //! output @param draftTokens, positionIds includes the lookahead and the verification branch information. - //! output @param draftLengthPtr holds the draft tokens length. - //! output @param attentionMask holds the draft tokens dependency mask, and attentionMaskOffset is the index offset - //! in attentionMask. - void prepare(TensorPtr const& draftTokens, TensorPtr const& positionIds, TensorPtr const& draftLengthPtr, - TensorPtr const& attentionMask, runtime::SizeType32 attentionMaskOffset, - TensorConstPtr const& lastPositionIdPtr, TensorConstPtr const& lastTokenPtr); - - //! @brief update the internal states and generate accepted tokens from @param outputTokens. - //! input @param sampledTokens is the all the tokens from the language model. - //! input @param endToken is the end token for `verify` early quit. - //! output @param acceptedTokens, acceptedOffsets in @param acceptedLength. - void update(TensorPtr const& acceptedTokens, TensorPtr const& acceptedOffsets, TensorPtr const& acceptedLength, - TensorConstPtr const& sampledTokens, TensorConstPtr const& endToken); - - //! generate attention @param mask from @param posIds. - static void posIdsToMask(TensorPtr const& mask, TensorConstPtr const& posIds); - - //! inplace encode the @param tokens and @param posIds according to attention @param masks, and record the offsets - //! in @param encodeMap. - static runtime::SizeType32 treeEncode( - TensorPtr const& tokens, TensorPtr const& posIds, TensorPtr const& masks, TensorPtr const& encodeMap); - -private: - //! @brief generate lookahead branch information. - //! input @param startPosId is the first position id of the draftTokens. - //! output @param draftTokens, positionIds of the lookahead branch. - //! @return the actual filled lookahead length. - runtime::SizeType32 lookahead( - TensorPtr const& draftTokens, TensorPtr const& positionIds, runtime::SizeType32 startPosId); - - //! @brief generate verification branch information. Also save the guessed tokens for future verification. - //! input @param startPosId the first position id. - //! input @param lastToken the last golden token for searching in the pool. - //! output @param guessTokens, guessIds of the verification branch. - //! @return the actual filled guess length. - runtime::SizeType32 guess(TensorPtr const& guessTokens, TensorPtr const& guessIds, runtime::SizeType32 startPosId, - runtime::TokenIdType lastToken); - - //! @brief verify the guessed tokens results and generate the longest accepted tokens. - //! input @param newLastToken is the new-generated last golden token. - //! input @param sampledTokens is the generated token results from the language model. - //! input @param endToken is the end token for early quit detection. - //! output @param accepted in @param acceptedLength, including the first golden one. - //! output @param acceptedOffsets is the offsets of draft tokens, excluding the first golden one. - void verify(TensorPtr const& accepted, TensorPtr const& acceptedOffsets, TensorPtr const& acceptedLength, - runtime::TokenIdType newLastToken, TensorConstPtr const& sampledTokens, TensorConstPtr const& endToken); - -private: - LookaheadPoolManager mPoolManager; - //! the random prefill tokens, - TensorPtr mPrefillsMax; // shape [mMaxN-2] - TensorPtr mPrefills; // shape [mN-2] - //! the look ahead branch window - TensorPtr mPastTokensMax; // shape [mMaxW * (mMaxN-1)] - TensorPtr mPastTokens; // shape [mW, (mN-1)] - //! the shifted mPastTokens as key tokens; - TensorPtr mKeyTokensMax; // shape [mMaxW] - TensorPtr mKeyTokens; // shape [mW] - //! all the moving tail golden tokens - TensorPtr mGoldenTokensMax; // shape[mMaxN*2-1] - TensorPtr mGoldenTokens; // shape[mN*2-1] - //! the same guess tokens from `guess` and used in `verify` - TensorPtr mGuessTokensMax; // shape [mMaxG*(mMaxN-1)] - TensorPtr mGuessTokens; // shape [mG*(mN-1)] - TensorPtr mDraftTokensMax; - TensorPtr mDraftTokens; - TensorPtr mAttentionMask; - TensorPtr mEncodeMapMax; - TensorPtr mEncodeMap; - TensorPtr mSampledTokensMax; - TensorPtr mSampledTokens; - - //! look ahead algorithm parameters, Window size, Level and Guess set size. - //! max for reserving resources and current for current request. - runtime::SizeType32 const mMaxW{0}; - runtime::SizeType32 const mMaxN{0}; - runtime::SizeType32 const mMaxG{0}; - runtime::SizeType32 mW{0}; - runtime::SizeType32 mN{0}; - runtime::SizeType32 mG{0}; - runtime::SizeType32 mRuntimeMaxDraftLen{0}; - runtime::SizeType32 mRuntimeMaxDraftPathLen{0}; - //! in prefilling mode when mFilling < mN-1. - runtime::SizeType32 mFilling; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp b/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp deleted file mode 100644 index 986f0e0b978e..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp +++ /dev/null @@ -1,469 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "lookaheadDecodingLayer.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/lookaheadAlgorithm.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/lookaheadModule.h" -#include -#include -#include - -namespace tensorrt_llm::layers -{ - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -template -LookaheadDecodingLayer::CpuAlgorithmResources::CpuAlgorithmResources(DecoderDomain const& decoderDomain) -{ - auto const maxBatchSize = decoderDomain.getBatchSize(); - auto const beamWidth = decoderDomain.getBeamWidth(); - auto const decodingTokens = decoderDomain.getMaxDecodingTokens(); - auto lookaheadModule - = std::dynamic_pointer_cast(decoderDomain.getSpeculativeDecodingModule()); - auto const [maxW, maxN, maxG] = lookaheadModule->getExecutionConfig().get(); - SizeType32 maxTokensPerStep, maxNumNewTokens, maxDraftLen, maxAcceptedDraftLen; - std::tie(maxTokensPerStep, maxNumNewTokens, maxDraftLen, maxAcceptedDraftLen) - = executor::LookaheadDecodingConfig(maxW, maxN, maxG).calculateSpeculativeResource(); - TLLM_CHECK_WITH_INFO(beamWidth == 1, "Lookahead requires beam width = 1"); - TLLM_CHECK_WITH_INFO(maxTokensPerStep == decodingTokens, "%d != %d", maxTokensPerStep, decodingTokens); - - for (SizeType32 id = 0; id < maxBatchSize; id++) - { - mAlgos.emplace_back(maxW, maxN, maxG, id); - } - - mPrompts.reserve(maxBatchSize); - for (auto bi = 0; bi < maxBatchSize; bi++) - { - mPrompts.emplace_back(BufferManager::cpu(ITensor::makeShape({0}), tensorrt_llm::DataType::kINT32)); - } - - auto const maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - mBatchSlots = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mTargetTokens - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mTokensPerStep = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mEndIds = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - - mOutputIds - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxNumNewTokens}), tensorrt_llm::DataType::kINT32); - mNewTokens = BufferManager::cpu( - ITensor::makeShape({maxTokensPerStep, maxBatchSize, beamWidth}), tensorrt_llm::DataType::kINT32); - mPathsOffsets - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); - mPathsOffsetsBatch - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); - mNumNewTokens = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mNumNewTokensCumSum = BufferManager::cpu(ITensor::makeShape({maxBatchSize + 1}), tensorrt_llm::DataType::kINT32); - mNextDraftTokens - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), tensorrt_llm::DataType::kINT32); - mNextDraftPosIds - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), tensorrt_llm::DataType::kINT32); - mGenerationLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mPositionOffsets - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mPositionIds - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mAttentionMask - = BufferManager::cpu(ITensor::makeShape({maxTokensPerStep, maxTokensPerStep}), tensorrt_llm::DataType::kBOOL); - mPackedMask = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep, - static_cast(divUp(maxTokensPerStep, 32))}), - tensorrt_llm::DataType::kINT32); - mNextDraftLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mSequenceLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); -} - -template -LookaheadDecodingLayer::LookaheadDecodingLayer( - DecoderDomain const& decoderDomain, std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mCpuAlgo(std::make_optional(decoderDomain)) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto lookaheadModule - = std::dynamic_pointer_cast(decoderDomain.getSpeculativeDecodingModule()); - - auto const maxBatchSize = mDecoderDomain.getBatchSize(); - auto const maxTokensPerStep = mDecoderDomain.getMaxDecodingTokens(); - auto const vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - auto const maxTopK = 1; - auto const maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - auto const maxBatchShape2D = ITensor::makeShape({maxBatchSize, maxTokensPerStep}); - - mWorkspaceSize = getTopKWorkspaceSize(maxBatchSize, maxTokensPerStep, maxTopK, vocabSizePadded); - mTargetTokensDevice = mBufferManager->gpu(maxBatchShape2D, tensorrt_llm::DataType::kINT32); - mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); - - mSetupWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(maxBatchShape1D, tensorrt_llm::DataType::kINT64)); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void LookaheadDecodingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(LookaheadDecodingLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - if (mCpuAlgo) - { - auto& algoConfigs = setupParams->algoConfigs; - TLLM_CHECK_WITH_INFO(algoConfigs.size() == 1 || algoConfigs.size() == static_cast(batchSize), - "Lookahead runtime configuration size should be either 1 or batchSize"); - - for (auto bi = 0; bi < batchSize; bi++) - { - PRINT_SHAPE(setupParams->prompt[bi]); - PRINT_TOKENS(setupParams->prompt[bi]); - mCpuAlgo->mPrompts[bi]->reshape(setupParams->prompt[bi]->getShape()); - mBufferManager->copy(*setupParams->prompt[bi], *mCpuAlgo->mPrompts[bi]); - } - - mBufferManager->getStream().synchronize(); // sync prompt gpu to cpu - - auto const batchSlotsRange = BufferRange(*batchSlots); - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - auto const gbi = batchSlotsRange[bi]; - SizeType32 bi1orN = (algoConfigs.size() == 1) ? 0 : bi; - TLLM_LOG_DEBUG("CPU ALGO [ %d ] setup prompt %s", gbi, D(mCpuAlgo->mPrompts[bi]).values().c_str()); - auto [w, n, g] = algoConfigs[bi1orN].get(); - SizeType32 runtimeTokensPerStep = 0; - std::tie(runtimeTokensPerStep, std::ignore, std::ignore, std::ignore) - = executor::LookaheadDecodingConfig(w, n, g).calculateSpeculativeResource(); - TLLM_CHECK_WITH_INFO(runtimeTokensPerStep <= mDecoderDomain.getMaxDecodingTokens(), - "runtime w(%d) n(%d) g(%d) exceeds maxTokensPerStep(%d)", w, n, g, - mDecoderDomain.getMaxDecodingTokens()); - PRINT_VALUES(mCpuAlgo->mPrompts[bi]); - auto seed = DefaultDecodingParams::getSeed(); - if (setupParams->randomSeed) - { - auto& seeds = setupParams->randomSeed.value(); - seed = seeds.size() == 1 ? seeds[0] : seeds[bi]; - } - mCpuAlgo->mAlgos[gbi].setup(mCpuAlgo->mPrompts[bi], w, n, g, seed); - } - - for (runtime::SizeType32 bi = 0; bi < batchSize; bi++) - { - SizeType32 gbi = batchSlotsRange[bi]; - (BufferRange(*mCpuAlgo->mGenerationLengths))[gbi] = 1; - (BufferRange(*mCpuAlgo->mNextDraftLengths))[gbi] = 0; - BufferLocation(*mCpuAlgo->mPositionOffsets).at(gbi, 0) = 0; - BufferRange packedMaskRange(*ITensor::at(mCpuAlgo->mPackedMask, {gbi})); - for (auto& mask : packedMaskRange) - { - mask = 0; - } - packedMaskRange[0] = 1; - - PRINT_SHAPE(mCpuAlgo->mGenerationLengths); - PRINT_SHAPE(setupParams->generationLengths); - PRINT_SHAPE(mCpuAlgo->mPositionOffsets); - PRINT_SHAPE(setupParams->positionOffsets); - PRINT_SHAPE(mCpuAlgo->mPackedMask); - PRINT_SHAPE(setupParams->attentionPackedMasks); - mBufferManager->copy( - *ITensor::at(mCpuAlgo->mGenerationLengths, {gbi}), *ITensor::at(setupParams->generationLengths, {gbi})); - mBufferManager->copy( - *ITensor::at(mCpuAlgo->mPositionOffsets, {gbi}), *ITensor::at(setupParams->positionOffsets, {gbi})); - mBufferManager->copy( - *ITensor::at(mCpuAlgo->mPackedMask, {gbi}), *ITensor::at(setupParams->attentionPackedMasks, {gbi})); - } - - mBufferManager->getStream().synchronize(); // sync outputs cpu to gpu - } - - workspace->initializeDeviceCurandStates( - setupParams->randomSeed, batchSize, workspace->getDeviceBatchSlots(), mCurandStatesDevice); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void LookaheadDecodingLayer::forwardAsync(std::shared_ptr const& outputParams, - std::shared_ptr const& inputParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(LookaheadDecodingLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(inputParams); - auto outputs = std::dynamic_pointer_cast(outputParams); - auto batchSize = inputs->localBatchSize; - - TLLM_CHECK_WITH_INFO(inputs->batchSlots, "Batch slots must be provided for LookaheadDecoding"); - TLLM_CHECK_WITH_INFO(inputs->curTokensPerStep, "curTokensPerStep must be provided for LookaheadDecoding"); - TLLM_CHECK_WITH_INFO(outputs->sequenceLength, "sequenceLength must be provided for LookaheadDecoding"); - TLLM_CHECK_WITH_INFO(inputs->logits, "logits must be provided for LookaheadDecoding"); - TLLM_CHECK_WITH_INFO(inputs->localBatchSize > 0, "batchSize must be"); - - TopKSamplingKernelParams params; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.batchSize = batchSize; - params.maxTopK = 1; - params.returnAllSelectedTokens = true; - params.maxTokensPerStep = mDecoderDomain.getMaxDecodingTokens(); - params.maxSeqLen = mDecoderDomain.getMaxDecodingTokens(); - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.logProbs = bufferCastOrNull(inputs->logits); - params.outputIds = bufferCast(*mTargetTokensDevice); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.curandState = reinterpret_cast(bufferCast(*mCurandStatesDevice)); - params.tokensPerStep = bufferCast(*inputs->curTokensPerStep.value()); - - TLLM_LOG_DEBUG( - "invokeBatchTopKSampling: maxBatchSize=%d, batchSize=%d, maxTopK=%d, maxTokensPerStep=%d, maxSeqLen=%d, " - "vocabSizePadded=%d", - params.maxBatchSize, params.batchSize, params.maxTopK, params.maxTokensPerStep, params.maxSeqLen, - params.vocabSizePadded); - - // Sample multiple tokens per request and store them to separate to be accepted/rejected later - // Sequence length is not modified, endIds is not checked, outputLogProbs are not supported. - // Finished state is not set. - invokeBatchTopKSampling(params, getStream()); - - if (mCpuAlgo) - { - forwardSyncCPU(outputs, inputs); - mGlobalSteps += 1; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t LookaheadDecodingLayer::getWorkspaceSize() const noexcept -{ - return std::max(mWorkspaceSize, mSetupWorkspaceSize); -} - -inline void initAttentionMask(TensorPtr const& mask, std::shared_ptr& bufferManager) -{ - bufferManager->setZero(*mask); - BufferLocation maskLocation(*mask); - auto maskShape = mask->getShape(); - for (SizeType32 i = 0; i < maskShape.d[0]; i++) - { - maskLocation.at(i, 0) = true; - } -} - -inline void convertBoolToInt32(TensorPtr const& dst, TensorConstPtr const& src) -{ - auto dstShape = dst->getShape(); - auto srcShape = src->getShape(); - TLLM_CHECK(dstShape.d[0] == srcShape.d[0]); - TLLM_CHECK(dstShape.d[1] * 32 >= srcShape.d[1]); - BufferLocation dstLocation(*dst); - BufferLocation srcLocation(*src); - - auto setBit = [](SizeType32& x, SizeType32 idx, bool value) { x |= (value << idx); }; - for (auto i = 0; i < srcShape.d[0]; i++) - { - for (auto j = 0; j < srcShape.d[1]; j++) - { - setBit(dstLocation.at(i, j / 32), j % 32, srcLocation.at(i, j)); - } - } -} - -template -void LookaheadDecodingLayer::forwardSyncCPU( - std::shared_ptr const& outputs, std::shared_ptr const& inputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(LookaheadDecodingLayer_forwardSyncCPU); - - mCpuAlgo->mBatchSlots->reshape(inputs->batchSlots->getShape()); - mBufferManager->copy(*inputs->batchSlots, *mCpuAlgo->mBatchSlots); - mBufferManager->copy(*inputs->curTokensPerStep.value(), *mCpuAlgo->mTokensPerStep); - mBufferManager->copy(*inputs->endIds, *mCpuAlgo->mEndIds); - mBufferManager->copy(*outputs->sequenceLength.value(), *mCpuAlgo->mSequenceLengths); - - mBufferManager->copy(*mTargetTokensDevice, *mCpuAlgo->mTargetTokens); - - if (outputs->prevDraftLengths) - { - mBufferManager->copy(*mCpuAlgo->mNextDraftLengths, *outputs->prevDraftLengths); - } - - mBufferManager->getStream().synchronize(); - - auto const batchSize = inputs->localBatchSize; - - BufferRange tokensPerStepRange(*mCpuAlgo->mTokensPerStep); - BufferRange endIdsRange(*mCpuAlgo->mEndIds); - BufferLocation newTokensLocation(*mCpuAlgo->mNewTokens); - BufferRange numNewTokensRange(*mCpuAlgo->mNumNewTokens); - BufferRange numNewTokensCumSumRange(*mCpuAlgo->mNumNewTokensCumSum); - BufferRange batchSlotsRange(*mCpuAlgo->mBatchSlots); - BufferRange generationLengthsRange(*mCpuAlgo->mGenerationLengths); - BufferRange nextDraftLengthsRange(*mCpuAlgo->mNextDraftLengths); - BufferRange sequenceLengthsRange(*mCpuAlgo->mSequenceLengths); - BufferLocation pathsOffsetLocation(*mCpuAlgo->mPathsOffsets); - BufferLocation pathsOffsetBatchLocation(*mCpuAlgo->mPathsOffsetsBatch); - BufferLocation outputIdsLocation(*mCpuAlgo->mOutputIds); - - mBufferManager->setZero(*mCpuAlgo->mPathsOffsets); - mBufferManager->setZero(*mCpuAlgo->mNumNewTokens); - mBufferManager->setZero(*mCpuAlgo->mNumNewTokensCumSum); - mBufferManager->setZero(*mCpuAlgo->mPackedMask); - - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - SizeType32 gbi = batchSlotsRange[bi]; - LookaheadAlgorithm& theAlgo(mCpuAlgo->mAlgos[gbi]); - - SizeType32 const tokensPerStep = generationLengthsRange[gbi]; - TensorPtr sampledTokens = ITensor::slice(mCpuAlgo->mTargetTokens, {gbi, 0}, tokensPerStep); - - if (tokensPerStep == 1) - { - // The first step in generation phase has no draft tokens. - theAlgo.accept(sampledTokens); - mBufferManager->copy(*sampledTokens, *ITensor::slice(mCpuAlgo->mOutputIds, {gbi, 0}, tokensPerStep)); - numNewTokensRange[gbi] = tokensPerStep; - BufferLocation(*mCpuAlgo->mNextDraftLengths).at(gbi) = 0; - } - else - { - theAlgo.update( // - ITensor::at(mCpuAlgo->mOutputIds, {gbi}), // - ITensor::at(mCpuAlgo->mPathsOffsets, {gbi}), // - ITensor::at(mCpuAlgo->mNumNewTokens, {gbi}), // - sampledTokens, // - ITensor::at(mCpuAlgo->mEndIds, {gbi})); - } - - auto maxNumNewTokens = mCpuAlgo->mOutputIds->getShape().d[1]; - - mBufferManager->copy(*ITensor::at(mCpuAlgo->mOutputIds, {gbi}), - *ITensor::slice(outputs->outputIds, {gbi, 0, sequenceLengthsRange[gbi]}, maxNumNewTokens)); - - sequenceLengthsRange[gbi] += numNewTokensRange[gbi]; - - initAttentionMask(mCpuAlgo->mAttentionMask, mBufferManager); - - theAlgo.prepare( // - ITensor::at(mCpuAlgo->mNextDraftTokens, {gbi}), // - ITensor::at(mCpuAlgo->mNextDraftPosIds, {gbi}), // - ITensor::at(mCpuAlgo->mNextDraftLengths, {gbi}), // - mCpuAlgo->mAttentionMask, 1, // - ITensor::at(mCpuAlgo->mSequenceLengths, {gbi}), // - ITensor::at(mCpuAlgo->mOutputIds, {gbi, numNewTokensRange[gbi] - 1})); - - convertBoolToInt32(ITensor::at(mCpuAlgo->mPackedMask, {gbi}), mCpuAlgo->mAttentionMask); - - BufferLocation posIdsLocation(*ITensor::at(mCpuAlgo->mPositionIds, {gbi})); - for (auto& posid : posIdsLocation) - { - posid = sequenceLengthsRange[gbi] - 1; - } - mBufferManager->copy(*ITensor::slice(mCpuAlgo->mNextDraftPosIds, {gbi, 0}, nextDraftLengthsRange[gbi]), - *ITensor::slice(mCpuAlgo->mPositionIds, {gbi, 1}, nextDraftLengthsRange[gbi])); - - BufferRange offsetRange(*ITensor::at(mCpuAlgo->mPositionOffsets, {gbi})); - for (size_t i = 0; i < posIdsLocation.size(); i++) - { - offsetRange[i] = posIdsLocation[i] - posIdsLocation[0]; - } - - TensorPtr accepted = ITensor::slice(mCpuAlgo->mOutputIds, {gbi, 0}, numNewTokensRange[gbi]); - TensorPtr draft = ITensor::slice(mCpuAlgo->mNextDraftTokens, {gbi, 0}, nextDraftLengthsRange[gbi]); - TLLM_LOG_DEBUG("CPU ALGO [ %d ] forward, %s", gbi, D(sampledTokens).values().c_str()); - TLLM_LOG_DEBUG("[%d][%d] CPU ALGO [ %d ] forward, %s, %s", mGlobalSteps, batchSize, gbi, - D(accepted).values().c_str(), D(draft).values().c_str()); - } - - size_t pi = 0; - numNewTokensCumSumRange[0] = 0; - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - SizeType32 gbi = batchSlotsRange[bi]; - SizeType32 acceptedDraftLen = numNewTokensRange[gbi] <= 1 ? 0 : (numNewTokensRange[gbi] - 1); - numNewTokensCumSumRange[bi + 1] = numNewTokensCumSumRange[bi] + acceptedDraftLen; - for (SizeType32 tj = 0; tj < acceptedDraftLen; tj++) - { - pathsOffsetBatchLocation[pi++] = pathsOffsetLocation.at(gbi, tj); - } - } - - while (pi < pathsOffsetBatchLocation.size()) - { - pathsOffsetBatchLocation[pi++] = 0; - } - - TLLM_CHECK(outputs->numNewTokens); - - mBufferManager->copy(*mCpuAlgo->mSequenceLengths, *outputs->sequenceLength.value()); - mBufferManager->copy(*mCpuAlgo->mNewTokens, *outputs->newTokens); - - mBufferManager->copy(*mCpuAlgo->mNumNewTokens, *outputs->numNewTokens.value()); - mBufferManager->copy(*mCpuAlgo->mPathsOffsetsBatch, *outputs->pathsOffsets); - mBufferManager->copy(*mCpuAlgo->mNumNewTokensCumSum, *outputs->numNewTokensCumSum); // - mBufferManager->copy(*mCpuAlgo->mNextDraftTokens, *outputs->nextDraftTokens); - - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - SizeType32 gbi = batchSlotsRange[bi]; - // nextDraftLengthsRange[gbi] = mDecoderDomain.getMaxDecodingTokens() - 1; - generationLengthsRange[gbi] = nextDraftLengthsRange[gbi] + 1; - } - - if (outputs->nextDraftLengths) - { - mBufferManager->copy(*mCpuAlgo->mNextDraftLengths, *outputs->nextDraftLengths); - } - - mBufferManager->copy(*mCpuAlgo->mPackedMask, *outputs->packedMasks); - mBufferManager->copy(*mCpuAlgo->mGenerationLengths, *outputs->generationLengths); - mBufferManager->copy(*mCpuAlgo->mPositionOffsets, *outputs->positionOffsets); - mBufferManager->copy(*mCpuAlgo->mPositionIds, *outputs->positionIds); - - mBufferManager->getStream().synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} -template class LookaheadDecodingLayer; -template class LookaheadDecodingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.h b/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.h deleted file mode 100644 index e20b59b22b6a..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.h +++ /dev/null @@ -1,94 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "lookaheadAlgorithm.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" - -namespace tensorrt_llm::layers -{ - -//! \brief LookaheadDecodingLayer -template -class LookaheadDecodingLayer : public BaseLayer -{ -public: - using Base = BaseLayer; - using Base::mBufferManager; - - LookaheadDecodingLayer(DecoderDomain const& decoderDomain, std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputParams, - std::shared_ptr const& inputParams, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - void forwardSyncCPU(std::shared_ptr const& outputs, - std::shared_ptr const& inputs); - -private: - using Base::mDecoderDomain; - - size_t mWorkspaceSize{}; - size_t mSetupWorkspaceSize{}; - TensorPtr mCurandStatesDevice; - TensorPtr mTargetTokensDevice; - - struct CpuAlgorithmResources - { - explicit CpuAlgorithmResources(DecoderDomain const& decoderDomain); - - std::vector mAlgos; - std::vector mPrompts; - TensorPtr mBatchSlots; - TensorPtr mTargetTokens; - TensorPtr mTokensPerStep; - TensorPtr mEndIds; - - TensorPtr mOutputIds; - TensorPtr mPathsOffsets; - TensorPtr mPathsOffsetsBatch; - TensorPtr mNumNewTokens; - TensorPtr mNumNewTokensCumSum; - TensorPtr mNewTokens; - - TensorPtr mNextDraftTokens; - TensorPtr mNextDraftPosIds; - TensorPtr mNextDraftLengths; - TensorPtr mSequenceLengths; - TensorPtr mGenerationLengths; - TensorPtr mAttentionMask; - TensorPtr mPackedMask; - TensorPtr mPositionOffsets; - TensorPtr mPositionIds; - }; - - std::optional mCpuAlgo; - - runtime::SizeType32 mGlobalSteps{0}; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h b/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h deleted file mode 100644 index 8e3e8f6c590d..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h +++ /dev/null @@ -1,433 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::layers -{ - -template -class BufferLocation : public runtime::BufferRange -{ -public: - using typename runtime::BufferRange::size_type; - using runtime::BufferRange::begin; - using runtime::BufferRange::operator[]; - - BufferLocation(T* data, size_type size) - : runtime::BufferRange{data, size} - { - } - - template , bool> = true> - explicit BufferLocation(runtime::ITensor& tensor) - : BufferLocation(runtime::bufferCast(tensor), tensor.getSize()) - { - mStrides = runtime::ITensor::strides(tensor.getShape()); - } - - template , bool> = true> - explicit BufferLocation(runtime::ITensor const& tensor) - : BufferLocation(runtime::bufferCast(tensor), tensor.getSize()) - { - mStrides = runtime::ITensor::strides(tensor.getShape()); - } - - inline T& at(runtime::ITensor::Shape const& dims) - { - return *ptr(dims); - } - - inline T& at(std::initializer_list const& dims) - { - return *ptr(dims); - } - - template - inline T& at(Args... args) - { - runtime::ITensor::DimType64 offset = 0; - runtime::ITensor::DimType64 dims = 0; - atHelper(offset, dims, args...); - return *(begin() + offset); - } - - inline T& operator[](runtime::ITensor::Shape const& dims) - { - return *ptr(dims); - } - - inline T& operator[](std::initializer_list const& dims) - { - return *ptr(dims); - } - - inline T* ptr(runtime::ITensor::Shape const& dims) - { - return begin() + offset(dims); - } - - inline T* ptr(std::initializer_list const& dims) - { - return ptr(runtime::ITensor::makeShape(dims)); - } - - runtime::ITensor::DimType64 offset(runtime::ITensor::Shape const& dims) - { - TLLM_CHECK(mStrides.nbDims == dims.nbDims); - runtime::ITensor::DimType64 result = 0; - for (runtime::ITensor::DimType64 di = 0; di < mStrides.nbDims; di++) - { - result += dims.d[di] * mStrides.d[di]; - } - return result; - } - - runtime::ITensor::DimType64 offset(std::initializer_list const& dims) - { - return offset(runtime::ITensor::makeShape(dims)); - } - -private: - inline void atHelper(runtime::ITensor::DimType64& offset, runtime::ITensor::DimType64& dims) {} - - template - inline void atHelper(runtime::ITensor::DimType64& offset, runtime::ITensor::DimType64& dims, int dim, Args... args) - { - offset += dim * mStrides.d[dims++]; - atHelper(offset, dims, args...); - } - -private: - runtime::ITensor::Shape mStrides; -}; - -class DebugTensor -{ -public: - DebugTensor(runtime::ITensor const& tensor, char const* name, - std::shared_ptr bufferManager = nullptr, - std::shared_ptr stream = nullptr) - : mTensor(tensor) - , mName(name) - , mBufferManager(bufferManager) - , mStream(stream) - { - } - - DebugTensor(runtime::ITensor::SharedConstPtr tensor, char const* name, - std::shared_ptr bufferManager = nullptr, - std::shared_ptr stream = nullptr) - : DebugTensor(*tensor, name, bufferManager, stream) - { - } - - uint8_t const& u8(std::initializer_list const& dims) - { - return (BufferLocation(mTensor))[dims]; - } - - uint8_t const& u8(int32_t idx) - { - return (BufferLocation(mTensor))[idx]; - } - - int8_t const& i8(std::initializer_list const& dims) - { - return (BufferLocation(mTensor))[dims]; - } - - int8_t const& i8(int32_t idx) - { - return (BufferLocation(mTensor))[idx]; - } - - int32_t const& i32(std::initializer_list const& dims) - { - return (BufferLocation(mTensor))[dims]; - } - - int32_t const& i32(int32_t idx) - { - return (BufferLocation(mTensor))[idx]; - } - - int64_t const& i64(std::initializer_list const& dims) - { - return (BufferLocation(mTensor))[dims]; - } - - int64_t const& i64(int32_t idx) - { - return (BufferLocation(mTensor))[idx]; - } - - float const& f(std::initializer_list const& dims) - { - return (BufferLocation(mTensor))[dims]; - } - - float const& f(int32_t idx) - { - return (BufferLocation(mTensor))[idx]; - } - - runtime::BufferManager::ITensorPtr copyToHostOptional() - { - runtime::BufferManager::ITensorPtr hostPtr{nullptr}; - if (mTensor.getMemoryType() == runtime::MemoryType::kGPU) - { - auto theManager = mBufferManager - ? mBufferManager - : std::make_shared(mStream ? mStream : std::make_shared()); - hostPtr = theManager->copyFrom(mTensor, runtime::MemoryType::kCPU); - theManager->getStream().synchronize(); - } - return hostPtr; - } - - std::string string(void) - { - runtime::BufferManager::ITensorPtr hostPtr = copyToHostOptional(); - runtime::BufferRange range(hostPtr ? (*hostPtr) : mTensor); - std::string result(range.size(), '\0'); - std::copy(range.begin(), range.end(), result.begin()); - return result; - } - - std::string tokens(void) - { - using namespace tensorrt_llm::runtime; - std::ostringstream buf; - auto shape = mTensor.getShape(); - runtime::BufferManager::ITensorPtr hostPtr = copyToHostOptional(); - runtime::BufferRange tensorRange(hostPtr ? (*hostPtr) : mTensor); - - buf << mName << ": " << mTensor.getMemoryTypeName() << ',' << mTensor.getDataTypeName() << ',' << shape; - auto line = [&buf](TokenIdType const* array, SizeType32 size) - { - buf << '['; - for (SizeType32 i = 0; i < size; i++) - { - auto token = array[i]; - if (token >= ' ' && token <= '~') - { - buf << '\'' << static_cast(token) << '\''; - } - else - { - buf << token; - } - if (i != size - 1) - { - buf << ','; - } - } - buf << ']'; - }; - if (shape.nbDims == 0) - { - buf << "[]"; - } - else if (shape.nbDims == 1) - { - line(tensorRange.begin(), shape.d[0]); - } - else if (shape.nbDims == 2) - { - buf << '['; - for (runtime::SizeType32 i = 0; i < shape.d[0]; i++) - { - buf << "\n " << i << ": "; - line(tensorRange.begin() + i * shape.d[1], shape.d[1]); - } - buf << ']'; - } - else - { - buf << "Too Large to be printed"; - } - return buf.str(); - } - - template - std::string values(void) - { - using namespace tensorrt_llm::runtime; - std::ostringstream buf; - auto shape = mTensor.getShape(); - runtime::BufferManager::ITensorPtr hostPtr = copyToHostOptional(); - runtime::BufferRange tensorRange(hostPtr ? (*hostPtr) : mTensor); - - buf << mName << ": " << mTensor.getMemoryTypeName() << ',' << mTensor.getDataTypeName() << ',' << shape; - auto line = [&buf](T const* array, SizeType32 size) - { - buf << '['; - for (SizeType32 i = 0; i < size; i++) - { - buf << static_cast(array[i]); - if (i != size - 1) - { - buf << ','; - } - } - buf << ']'; - }; - if (shape.nbDims == 0) - { - buf << "[]"; - } - else if (shape.nbDims == 1) - { - line(tensorRange.begin(), shape.d[0]); - } - else if (shape.nbDims == 2) - { - buf << '['; - for (runtime::SizeType32 i = 0; i < shape.d[0]; i++) - { - buf << "\n " << i << ": "; - line(tensorRange.begin() + i * shape.d[1], shape.d[1]); - } - buf << ']'; - } - else - { - buf << "Too Large to be printed"; - } - return buf.str(); - } - - std::string values(void) - { - switch (mTensor.getDataType()) - { - case tensorrt_llm::DataType::kBOOL: return values(); - case tensorrt_llm::DataType::kFLOAT: return values(); - case tensorrt_llm::DataType::kINT8: return values(); - case tensorrt_llm::DataType::kINT32: return values(); - case tensorrt_llm::DataType::kINT64: return values(); - case tensorrt_llm::DataType::kUINT8: return values(); - default: return std::string(mName + ": Unsupported data type"); - } - } - - std::string shape(void) - { - using namespace tensorrt_llm::runtime; - std::ostringstream buf; - buf << mName << ": " << mTensor.getShape(); - return buf.str(); - } - - void print_tokens(void) - { - TLLM_LOG_DEBUG(tokens()); - } - - void print_values(void) - { - TLLM_LOG_DEBUG(values()); - } - - void print_shape(void) - { - TLLM_LOG_DEBUG(shape()); - } - - template - void randomize(runtime::SizeType32 vtype) - { - runtime::BufferRange tensorRange(const_cast(mTensor)); - for (auto& item : tensorRange) - { - item = vtype == 0 ? 0 : vtype == 1 ? 1 : rand(); - } - } - - void randomize(void) - { - if (mTensor.getMemoryType() == runtime::MemoryType::kGPU) - { - runtime::ITensor& nonConstTensor = const_cast(mTensor); - runtime::BufferManager manager{std::make_shared()}; - runtime::ITensor::SharedConstPtr cpuBuffer = manager.cpu(mTensor.getShape(), mTensor.getDataType()); - DebugTensor(cpuBuffer, "cpuBuffer").randomize(); - manager.copy(*cpuBuffer, nonConstTensor); - manager.getStream().synchronize(); - } - else - { - switch (mTensor.getDataType()) - { - case tensorrt_llm::DataType::kBOOL: return randomize(3); - case tensorrt_llm::DataType::kFLOAT: return randomize(3); - case tensorrt_llm::DataType::kINT8: return randomize(3); - case tensorrt_llm::DataType::kINT32: return randomize(3); - case tensorrt_llm::DataType::kINT64: return randomize(3); - case tensorrt_llm::DataType::kUINT8: return randomize(3); - default: return; - } - } - } - - void setZeros(void) - { - switch (mTensor.getDataType()) - { - case tensorrt_llm::DataType::kBOOL: return randomize(0); - case tensorrt_llm::DataType::kFLOAT: return randomize(0); - case tensorrt_llm::DataType::kINT8: return randomize(0); - case tensorrt_llm::DataType::kINT32: return randomize(0); - case tensorrt_llm::DataType::kINT64: return randomize(0); - case tensorrt_llm::DataType::kUINT8: return randomize(0); - default: return; - } - } - - void setOnes(void) - { - switch (mTensor.getDataType()) - { - case tensorrt_llm::DataType::kBOOL: return randomize(1); - case tensorrt_llm::DataType::kFLOAT: return randomize(1); - case tensorrt_llm::DataType::kINT8: return randomize(1); - case tensorrt_llm::DataType::kINT32: return randomize(1); - case tensorrt_llm::DataType::kINT64: return randomize(1); - case tensorrt_llm::DataType::kUINT8: return randomize(1); - default: return; - } - } - -private: - runtime::ITensor const& mTensor; - std::string mName; - std::shared_ptr mBufferManager; - std::shared_ptr mStream; -}; - -#define D(x) tensorrt_llm::layers::DebugTensor(x, #x) -#define Db(x, bufferManager) tensorrt_llm::layers::DebugTensor(x, #x, bufferManager, nullptr) -#define Ds(x, stream) tensorrt_llm::layers::DebugTensor(x, #x, nullptr, stream) -#define PRINT_TOKENS(x) D(x).print_tokens() -#define PRINT_VALUES(x) D(x).print_values() -#define PRINT_SHAPE(x) D(x).print_shape() - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp b/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp deleted file mode 100644 index 397b4262226a..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp +++ /dev/null @@ -1,119 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/layers/lookaheadPoolManager.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include - -namespace tensorrt_llm::layers -{ - -using namespace tensorrt_llm::runtime; - -void LookaheadPoolManager::setup(SizeType32 guessSetSize) -{ - TLLM_CHECK(guessSetSize >= 0 && guessSetSize <= mGuessSetSizeMax); - mGuessSetSize = guessSetSize; - mTokenMap.clear(); -} - -void LookaheadPoolManager::insertOne(Key key, TensorConstPtr const& ngram) -{ - if (TLLM_UNLIKELY(ITensor::volume(ngram->getShape()) == 0 || mGuessSetSize == 0)) - { - return; - } - - auto search = mTokenMap.find(key); - if (search != mTokenMap.end()) - { - search->second.remove_if( - [&ngram](TensorConstPtr const& item) - { - BufferRange ngramRange(*ngram); - BufferRange itemRange(*item); - return std::equal(ngramRange.begin(), ngramRange.end(), itemRange.begin()); - }); - if (mGuessSetSize > 0 && search->second.size() >= static_cast(mGuessSetSize)) - { - search->second.pop_front(); - } - search->second.push_back(ngram); - } - else - { - mTokenMap.insert({key, std::list({ngram})}); - } -} - -void LookaheadPoolManager::accept(TensorConstPtr const& prompt, SizeType32 level) -{ - SizeType32 length = prompt->getShape().d[0]; - BufferRange promptRange(*prompt); - for (SizeType32 ti = 0; ti + level - 1 < length; ti++) - { - auto key = promptRange[ti]; - TensorPtr ngram = BufferManager::cpu(ITensor::makeShape({level - 1}), tensorrt_llm::DataType::kINT32); - BufferRange sourceRange(*ITensor::slice(prompt, ti + 1, level - 1)); - BufferRange ngramRange(*ngram); - std::copy(sourceRange.begin(), sourceRange.end(), ngramRange.begin()); - - insertOne(key, ngram); - } -} - -std::list LookaheadPoolManager::guess(Key lastToken, SizeType32 guessSize) const -{ - auto search = mTokenMap.find(lastToken); - if (search != mTokenMap.end()) - { - auto ngrams = search->second; - if (ngrams.size() > static_cast(guessSize)) - { - auto it = std::prev(ngrams.end(), guessSize); - return std::list(it, ngrams.end()); - } - else - { - return ngrams; - } - } - else - { - return std::list(); - } -} - -void LookaheadPoolManager::update(TensorConstPtr const& keyTokens, TensorConstPtr const& ngramTokens) -{ - TLLM_CHECK(keyTokens->getShape().d[0] == ngramTokens->getShape().d[0]); - BufferRange keyRange(*keyTokens); - auto window = ngramTokens->getShape().d[0]; - - for (SizeType32 wi = 0; wi < window; wi++) - { - TensorConstPtr source = ITensor::at(ngramTokens, {wi}); - TensorPtr ngram = BufferManager::cpu(source->getShape(), tensorrt_llm::DataType::kINT32); - BufferRange sourceRange(*source); - BufferRange ngramRange(*ngram); - std::copy(sourceRange.begin(), sourceRange.end(), ngramRange.begin()); - insertOne(keyRange[wi], ngram); - } -} - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadPoolManager.h b/cpp/tensorrt_llm/layers/lookaheadPoolManager.h deleted file mode 100644 index cc3bba694856..000000000000 --- a/cpp/tensorrt_llm/layers/lookaheadPoolManager.h +++ /dev/null @@ -1,77 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include -#include - -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::layers -{ - -//! @brief A helper class for managing key-ngram pool. -class LookaheadPoolManager -{ -public: - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorConstPtr = runtime::ITensor::SharedConstPtr; - using Key = runtime::TokenIdType; - - LookaheadPoolManager(runtime::SizeType32 maxG) - : mGuessSetSizeMax(maxG) - { - } - - //! @brief setup runtime resource - //! @param guessSetSize the runtime guessSetSize. - void setup(runtime::SizeType32 guessSetSize); - - //! @brief fill token map from accepted tokens, including prompt. - //! @param prompt the user input prompt, [length] on cpu - //! @param level the n-gram length - void accept(TensorConstPtr const& prompt, runtime::SizeType32 level); - - //! @brief get a list of guess tokens - //! @param lastToken the newest golden token - //! @param guessSize at most guessSize candidates returned - //! @return the list guess tokens, with list size <= guessSize - std::list guess(Key lastToken, runtime::SizeType32 guessSize) const; - - //! @brief update token map with new generated tokens - //! @param keyTokens the new shifted out tokens from each window, as the key, [window] on cpu - //! @param ngramTokens the new shifted lookahead window, as the ngrams, [window, ngramLen] on cpu - void update(TensorConstPtr const& keyTokens, TensorConstPtr const& ngramTokens); - - std::unordered_map> const& getMap() const - { - return mTokenMap; - } - -private: - void insertOne(Key key, TensorConstPtr const& ngram); - -private: - //! @brief the token map with token as key and list of n-gram as value - std::unordered_map> mTokenMap; - //! @brief guess set size, -1 for infinite size - runtime::SizeType32 const mGuessSetSizeMax; - runtime::SizeType32 mGuessSetSize; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp b/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp deleted file mode 100644 index 9e4098b34ebf..000000000000 --- a/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp +++ /dev/null @@ -1,483 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "medusaDecodingLayer.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -MedusaDecodingLayer::MedusaDecodingLayer( - DecoderDomain const& decoderDomain, std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - allocateBuffer(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const maxDraftPathLen = mDecoderDomain.getSpeculativeDecodingModule()->getMaxDraftPathLen(); - // Get sampling workspace size - { - auto samplingSizePrimarySampling = getTopKWorkspaceSize(mDecoderDomain.getBatchSize(), - mDecoderDomain.getMaxDecodingTokens(), TOP_K_MAX, mDecoderDomain.getVocabSizePadded()); - - auto const maxBatchSizeHeadNums = mDecoderDomain.getBatchSize() * maxDraftPathLen; - auto samplingSizeMedusaHeadsSampling - = getTopKWorkspaceSize(maxBatchSizeHeadNums, 1, TOP_K_MAX, mDecoderDomain.getVocabSizePadded()); - - mWorkspaceSize = std::max(samplingSizePrimarySampling, samplingSizeMedusaHeadsSampling); - } - - mDraftIdsPtrHost = BufferManager::pinnedPool( - ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize()), maxDraftPathLen}), - TRTDataType::value); - mCummulativeTopK.resize(mDecoderDomain.getBatchSize() * maxDraftPathLen); - - auto const batchSize = mDecoderDomain.getBatchSize(); - auto const batchSizeShape = ITensor::makeShape({mDecoderDomain.getBatchSize()}); - mCurandStatesDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(batchSize * sizeof(curandState_t))}), TRTDataType::value); - mRuntimeTopKDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mTargetTokensDevice = mBufferManager->gpu( - ITensor::makeShape({batchSize, mDecoderDomain.getMaxDecodingTokens()}), TRTDataType::value); - mRandomSeedsDevice - = mBufferManager->gpu(ITensor::makeShape({batchSize * maxDraftPathLen}), TRTDataType::value); - mMedusaSelectedLogitsPtrsDevice - = mBufferManager->gpu(ITensor::makeShape({batchSize, maxDraftPathLen}), TRTDataType::value); - mCurandStatesMedusaLogitsDevice = mBufferManager->gpu( - ITensor::makeShape({batchSize, maxDraftPathLen, sizeof(curandState_t)}), TRTDataType::value); - mRuntimeTopKPerRequestPerMedusaHeadDevice - = mBufferManager->gpu(ITensor::makeShape({batchSize, maxDraftPathLen}), TRTDataType::value); - mNewDraftTokensDevice = mBufferManager->gpu( - ITensor::makeShape({batchSize, mDecoderDomain.getMaxDecodingTokens()}), TRTDataType::value); - mBestPathIdsDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - - mTiledBatchSlotsSetup = BufferManager::pinnedPool( - ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - tensorrt_llm::DataType::kINT32); - mTiledBatchSlotsForward = BufferManager::pinnedPool( - ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - tensorrt_llm::DataType::kINT32); - mMedusaInputLogitsPtrs = BufferManager::pinnedPool( - ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - TRTDataType::value); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(MedusaDecodingLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - workspace->initializeDeviceCurandStates( - setupParams->randomSeed, batchSize, workspace->getDeviceBatchSlots(), mCurandStatesDevice); - - auto const maxDraftPathLen = mDecoderDomain.getSpeculativeDecodingModule()->getMaxDraftPathLen(); - auto const batchSizeMaxNumHeads = batchSize * maxDraftPathLen; - auto randomSeed = setupParams->randomSeed.value_or(std::vector(batchSize, uint64_t{0})); - std::vector tiledRandomSeed(batchSizeMaxNumHeads); - if (randomSeed.size() > 1) - { - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - tiledRandomSeed[bi * maxDraftPathLen + hi] = randomSeed[bi]; - } - } - } - auto* tiledBatchSlots = bufferCast(*mTiledBatchSlotsSetup); - BufferRange batchSlotsRange(*batchSlots); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - tiledBatchSlots[bi * maxDraftPathLen + hi] = batchSlotsRange[bi] + hi; - } - } - auto tiledRandomSeedOpt = std::make_optional(std::move(tiledRandomSeed)); - workspace->initializeDeviceCurandStates( - tiledRandomSeedOpt, batchSizeMaxNumHeads, mTiledBatchSlotsSetup, mCurandStatesMedusaLogitsDevice); - - // Prepare runtime top K - auto prepareRuntimeTopK = [this, workspace](std::vector const& runtimeTopK, SizeType32 batchSize, - BufferConstPtr const& batchSlots, BufferPtr const& runtimeTopKDevice) - { - TLLM_CHECK_WITH_INFO(runtimeTopK.size() == 1 || runtimeTopK.size() == static_cast(batchSize), - fmtstr("runtimeTopK.size() (%lu) == batchSize (%d) is not satisfied!", runtimeTopK.size(), batchSize)); - SizeType32* topKSetupPtr = nullptr; - if (runtimeTopK.size() > 1) - { - DecodingLayerWorkspace::copyToWorkspace( - *this->mBufferManager, runtimeTopK, workspace->getWorkspaceDeviceBuffer()); - topKSetupPtr = workspace->getWorkspaceDevicePtrAs(); - } - auto* runtimeTopKDevicePtr = bufferCastOrNull(runtimeTopKDevice); - auto const* batchSlotsPtr = bufferCastOrNull(batchSlots); - invokeScatterDecodingParams( - topKSetupPtr, runtimeTopK.front(), runtimeTopKDevicePtr, batchSlotsPtr, batchSize, getStream()); - - // FIXME: monotonically growing - auto const curMaxTopK = *std::max_element(std::begin(runtimeTopK), std::end(runtimeTopK)); - return curMaxTopK; - }; - - SizeType32 constexpr defaultTopK = 1; - { - auto runtimeTopK = setupParams->runtimeTopK.value_or(std::vector{defaultTopK}); - auto const curMaxTopK - = prepareRuntimeTopK(runtimeTopK, batchSize, workspace->getDeviceBatchSlots(), mRuntimeTopKDevice); - mRuntimeMaxTopK = std::max(mRuntimeMaxTopK, curMaxTopK); - } - { - auto runtimeHeadsTopK = setupParams->runtimeHeadsTopK; - std::vector runtimeHeadsTopKFlatten; - if (runtimeHeadsTopK.has_value() && static_cast(runtimeHeadsTopK->size())) - { - for (auto const& sub : runtimeHeadsTopK.value()) - { - runtimeHeadsTopKFlatten.insert(runtimeHeadsTopKFlatten.end(), sub.begin(), sub.end()); - } - } - else - { - runtimeHeadsTopKFlatten = std::vector(batchSizeMaxNumHeads, defaultTopK); - } - - BufferRange batchSlotsRange(*batchSlots); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto const slot = batchSlotsRange[bi]; - SizeType32 cummulativeTopK = 0; - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - mCummulativeTopK[slot * maxDraftPathLen + hi] = cummulativeTopK; - cummulativeTopK += runtimeHeadsTopKFlatten[bi * maxDraftPathLen + hi]; - } - } - - auto* tiledBatchSlots = bufferCast(*mTiledBatchSlotsSetup); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - tiledBatchSlots[bi * maxDraftPathLen + hi] = maxDraftPathLen * batchSlotsRange[bi] + hi; - } - } - - auto const curMaxTopK - = prepareRuntimeTopK(runtimeHeadsTopKFlatten, static_cast(batchSizeMaxNumHeads), - mTiledBatchSlotsSetup, mRuntimeTopKPerRequestPerMedusaHeadDevice); - mRuntimeMaxTopKPerRequestPerMedusaHead = std::max(mRuntimeMaxTopKPerRequestPerMedusaHead, curMaxTopK); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(MedusaDecodingLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - auto outputs = std::dynamic_pointer_cast(baseOutputs); - - // TODO add typical acceptance similarly to EagleSampleAndAcceptDraftTokensPlugin::doTypicalAcceptance. - samplePrimeHeadTokens(*outputs, *inputs, workspace); - - acceptDraftTokens(*outputs, *inputs, workspace); - - sampleNewDraftTokens(*outputs, *inputs, workspace); - - scatterNewDraftTokens(*outputs, *inputs); - - packAcceptedPaths(*outputs, *inputs, workspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t MedusaDecodingLayer::getWorkspaceSize() const noexcept -{ - return std::max(mWorkspaceSize, mSetupWorkspaceSize); -} - -template -void MedusaDecodingLayer::samplePrimeHeadTokens(SpeculativeDecodingOutputs const& outputs, - MedusaDecodingInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.logits.value()->getDimension<0>(); - - auto logits = bufferCast(*inputs.logits.value()); - auto const* batchSlots = workspace->getDeviceBatchSlotsPtr(); - auto* sequenceLengths = bufferCastOrNull(outputs.sequenceLength); - auto* tokensPerStepDevice = bufferCast(*inputs.curTokensPerStep.value()); - - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(sequenceLengths != nullptr, "Sequence lengths must be provided for MedusaDecoding"); - - TopKSamplingKernelParams params; - params.logProbs = logits; - params.outputIds = bufferCastOrNull(mTargetTokensDevice); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.maxTopK = mRuntimeMaxTopK; - params.topKs = bufferCastOrNull(mRuntimeTopKDevice); - params.batchSlots = batchSlots; - params.curandState = reinterpret_cast(bufferCastOrNull(mCurandStatesDevice)); - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.tokensPerStep = tokensPerStepDevice; - params.maxTokensPerStep = mDecoderDomain.getMaxDecodingTokens(); - params.maxSeqLen = mDecoderDomain.getMaxDecodingTokens(); - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - - // Sample multiple tokens per request and store them to separate to be accepted/rejected later - // Sequence length is not modified, endIds is not checked, outputLogProbs are not supported. - // Finished state is not set. - invokeBatchTopKSampling(params, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::acceptDraftTokens(SpeculativeDecodingOutputs const& outputs, - MedusaDecodingInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.logits.value()->getDimension<0>(); - auto const maxSeqLen = outputs.outputIds->getDimension<-1>(); - - auto* outputIds = bufferCast(*outputs.outputIds); - auto const* endIds = bufferCast(*inputs.endIds); - auto const* paths = bufferCast(*inputs.paths); - - auto const* batchSlots = bufferCast(*inputs.batchSlots); - auto* sequenceLengths = bufferCastOrNull(outputs.sequenceLength); - auto* numNewTokens = bufferCast(*outputs.numNewTokens.value()); - auto* curTokensPerStepDevice = bufferCast(*inputs.curTokensPerStep.value()); - auto const* targetTokensPerStepDevice = bufferCast(*inputs.targetTokensPerStep); - - auto const maxDraftPathLen = mDecoderDomain.getSpeculativeDecodingModule()->getMaxDraftPathLen(); - - auto medusaInputLogitsPtrs = BufferRange(*mMedusaInputLogitsPtrs); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto const slot = batchSlots[bi]; - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - medusaInputLogitsPtrs[slot * maxDraftPathLen + hi] = bufferCast(*inputs.medusaLogits[slot][hi]); - } - } - - auto* draftIds = bufferCast(*outputs.nextDraftTokens); - - TLLM_CHECK_WITH_INFO(draftIds != nullptr, "Draft ids must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(sequenceLengths != nullptr, "Sequence lengths must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(numNewTokens != nullptr, "Accepted lengths must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO( - curTokensPerStepDevice != nullptr, "Current tokens per step must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO( - targetTokensPerStepDevice != nullptr, "Target tokens per step must be provided for MedusaDecoding"); - - // Compare draft tokens from outputIds with sampled target tokens at mTargetTokensDevice using paths. - // Select the longest accepted path, modify outputIds in-place, increment sequenceLengths accordingly. - // Fill mMedusaSelectedLogitsPtrsDevice with respective Medusa logits - auto* targetTokensDevicePtr = bufferCast(*mTargetTokensDevice); - auto* finishedStatesPtr - = reinterpret_cast(bufferCastOrNull(outputs.finished)); - auto* bestPathIdsDevicePtr = bufferCastOrNull(mBestPathIdsDevice); - auto medusaInputLogitsPtrsPtr = reinterpret_cast(bufferCast(*mMedusaInputLogitsPtrs)); - auto medusaSelectedLogitsPtrsDevicePtr - = const_cast(bufferCastOrNull(mMedusaSelectedLogitsPtrsDevice)); - - AcceptDraftTokensByIdsWithPathsParams params; - params.outputIds = outputIds; - params.draftIds = draftIds; - params.targetIds = targetTokensDevicePtr; - params.sequenceLengths = sequenceLengths; - params.acceptedLengths = numNewTokens; - params.finishedFinal = finishedStatesPtr; - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.paths = paths; - params.endIds = endIds; - params.medusaLogits = medusaInputLogitsPtrsPtr; - params.logitsPtrs = medusaSelectedLogitsPtrsDevicePtr; - params.curTokensPerStep = curTokensPerStepDevice; - params.targetTokensPerStep = targetTokensPerStepDevice; - params.bestPathIds = bestPathIdsDevicePtr; - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.vocabSize = mDecoderDomain.getVocabSize(); - params.maxSeqLen = maxSeqLen; - params.maxDraftPathLen = maxDraftPathLen; - params.maxDecodingTokens = mDecoderDomain.getMaxDecodingTokens(); - params.stream = getStream(); - - params.checkParams(); - - acceptDraftTokensByIdsWithPaths(params); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::sampleNewDraftTokens(SpeculativeDecodingOutputs const& outputs, - MedusaDecodingInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.logits.value()->getDimension<0>(); - auto const* batchSlots = bufferCast(*inputs.batchSlots); - auto* sequenceLengths = bufferCastOrNull(outputs.sequenceLength); - - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(sequenceLengths != nullptr, "Sequence lengths must be provided for MedusaDecoding"); - - auto const maxDraftPathLen = mDecoderDomain.getSpeculativeDecodingModule()->getMaxDraftPathLen(); - // For each request we sample Head Num times for topK[hi] tokens - auto const batchSizeHeadNums = batchSize * maxDraftPathLen; - auto const maxBatchSizeHeadNums = mDecoderDomain.getBatchSize() * maxDraftPathLen; - - auto* tiledBatchSlots = bufferCast(*mTiledBatchSlotsForward); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - tiledBatchSlots[bi * maxDraftPathLen + hi] = maxDraftPathLen * batchSlots[bi] + hi; - } - } - - auto* draftIdsPtrs = reinterpret_cast(bufferCast(*mDraftIdsPtrHost)); - - auto* newDraftTokensDeviceRange = bufferCast(*mNewDraftTokensDevice); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto slot = batchSlots[bi]; - for (SizeType32 hi = 0; hi < maxDraftPathLen; ++hi) - { - draftIdsPtrs[slot * maxDraftPathLen + hi] = newDraftTokensDeviceRange - + slot * mDecoderDomain.getMaxDecodingTokens() + mCummulativeTopK[slot * maxDraftPathLen + hi]; - } - } - - TopKSamplingKernelParams params{}; - params.logProbsPtrs = bufferCastOrNull(mMedusaSelectedLogitsPtrsDevice); - params.outputIdsPtrs = draftIdsPtrs; - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.maxTopK = mRuntimeMaxTopKPerRequestPerMedusaHead; - params.topKs = bufferCastOrNull(mRuntimeTopKPerRequestPerMedusaHeadDevice); - params.batchSlots = tiledBatchSlots; - params.curandState = reinterpret_cast(bufferCastOrNull(mCurandStatesMedusaLogitsDevice)); - params.batchSize = batchSizeHeadNums; - params.maxBatchSize = maxBatchSizeHeadNums; - params.maxTokensPerStep = 1; - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - params.returnAllSelectedTokens = true; - - invokeBatchTopKSampling(params, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::scatterNewDraftTokens( - SpeculativeDecodingOutputs const& outputs, MedusaDecodingInputs const& inputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.logits.value()->getDimension<0>(); - auto const* batchSlots = bufferCast(*inputs.batchSlots); - - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for MedusaDecoding"); - - auto* draftIds = bufferCastOrNull(outputs.nextDraftTokens); - auto* tokensPerStepDevice = bufferCastOrNull(inputs.curTokensPerStep); - auto const* treeIds = bufferCastOrNull(inputs.treeIds); - TLLM_CHECK_WITH_INFO(draftIds != nullptr, "Draft ids must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(tokensPerStepDevice != nullptr, "Tokens per step must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(treeIds != nullptr, "Tree ids must be provided for MedusaDecoding"); - - auto* newDraftTokensDevice = bufferCastOrNull(mNewDraftTokensDevice); - scatterMedusaDraftTokens(draftIds, newDraftTokensDevice, treeIds, tokensPerStepDevice, batchSlots, - mDecoderDomain.getMaxDecodingTokens(), batchSize, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void MedusaDecodingLayer::packAcceptedPaths(SpeculativeDecodingOutputs const& outputs, - MedusaDecodingInputs const& inputs, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputs.logits.value()->getDimension<0>(); - auto const* paths = bufferCast(*inputs.paths); - auto const* batchSlots = workspace->getDeviceBatchSlotsPtr(); - auto* numNewTokens = bufferCast(*outputs.numNewTokens.value()); - auto* numNewTokensCumSum = bufferCast(*outputs.numNewTokensCumSum); - auto* pathsOffsets = bufferCast(*outputs.pathsOffsets); - auto* bestPathIdsDevicePtr = bufferCastOrNull(mBestPathIdsDevice); - - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(numNewTokens != nullptr, "Accepted lengths must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(numNewTokensCumSum != nullptr, "numNewTokensCumSum must be provided for MedusaDecoding"); - TLLM_CHECK_WITH_INFO(pathsOffsets != nullptr, "pathsOffsets must be provided for MedusaDecoding"); - invokePackAcceptedPaths(numNewTokensCumSum, pathsOffsets, numNewTokens, bestPathIdsDevicePtr, paths, batchSlots, - batchSize, batchSize, mDecoderDomain.getMaxDecodingTokens(), - mDecoderDomain.getSpeculativeDecodingModule()->getMaxPathLen(), false, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class MedusaDecodingLayer; -template class MedusaDecodingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/medusaDecodingLayer.h b/cpp/tensorrt_llm/layers/medusaDecodingLayer.h deleted file mode 100644 index cce424963aa4..000000000000 --- a/cpp/tensorrt_llm/layers/medusaDecodingLayer.h +++ /dev/null @@ -1,89 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include - -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" - -namespace tensorrt_llm::layers -{ - -//! \brief -template -class MedusaDecodingLayer : public BaseLayer -{ -public: - using Base = BaseLayer; - using PathsVec = std::vector>>; - - MedusaDecodingLayer(DecoderDomain const& decoderDomain, std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - void allocateBuffer(); - - void samplePrimeHeadTokens(SpeculativeDecodingOutputs const& outputs, MedusaDecodingInputs const& inputs, - std::shared_ptr const& workspace); - void acceptDraftTokens(SpeculativeDecodingOutputs const& outputs, MedusaDecodingInputs const& inputs, - std::shared_ptr const& workspace); - void sampleNewDraftTokens(SpeculativeDecodingOutputs const& outputs, MedusaDecodingInputs const& inputs, - std::shared_ptr const& workspace); - void scatterNewDraftTokens(SpeculativeDecodingOutputs const& outputs, MedusaDecodingInputs const& inputs); - void packAcceptedPaths(SpeculativeDecodingOutputs const& outputs, MedusaDecodingInputs const& inputs, - std::shared_ptr const& workspace); - -private: - using Base::mDecoderDomain; - - size_t mWorkspaceSize{0}; - size_t mSetupWorkspaceSize{0}; - runtime::SizeType32 mRuntimeMaxTopK{0}; - runtime::SizeType32 mRuntimeMaxTopKPerRequestPerMedusaHead{0}; - - TensorPtr mCurandStatesDevice; - TensorPtr mRuntimeTopKDevice; - TensorPtr mTargetTokensDevice; - TensorPtr mRandomSeedsDevice; - TensorPtr mMedusaSelectedLogitsPtrsDevice; - TensorPtr mCurandStatesMedusaLogitsDevice; - TensorPtr mRuntimeTopKPerRequestPerMedusaHeadDevice; - TensorPtr mNewDraftTokensDevice; - TensorPtr mBestPathIdsDevice; - - TensorPtr mTiledBatchSlotsSetup; - TensorPtr mTiledBatchSlotsForward; - TensorPtr mDraftIdsPtrHost; - TensorPtr mMedusaInputLogitsPtrs; - - std::vector mCummulativeTopK; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/penaltyLayer.cpp b/cpp/tensorrt_llm/layers/penaltyLayer.cpp deleted file mode 100644 index c72b8e463bc6..000000000000 --- a/cpp/tensorrt_llm/layers/penaltyLayer.cpp +++ /dev/null @@ -1,384 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "penaltyLayer.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/penaltyKernels.h" -#include "tensorrt_llm/kernels/penaltyTypes.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" - -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -size_t PenaltyLayer::getWorkspaceSize() const noexcept -{ - return mWorkspaceSize; -} - -template -PenaltyLayer::PenaltyLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - initialize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void PenaltyLayer::initialize() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - allocateBuffer(); - - mCyclicStep = 0; - mRuntimeMaxSeqLen = 0; - mConfiguredBeamWidth = -1; - - if (!mDecodingMode.isAuto()) - { - mConfiguredBeamWidth = mDecoderDomain.getBeamWidth(); - - allocateWorkspace(); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void PenaltyLayer::allocateWorkspace() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (mDecodingMode.isUseOccurrencePenalty()) - { - - auto const workspaceSize = mDecoderDomain.getBatchSize() * mDecoderDomain.getMaxDecodingTokens() - * mConfiguredBeamWidth * mDecoderDomain.getVocabSize() * 2; - mPenaltyWorkspaceDevice = mBufferManager->gpu(workspaceSize, tensorrt_llm::DataType::kINT32); - - if (mDecodingMode.isBeamSearch()) - { - mPenaltyWorkspacePrevDevice = mBufferManager->gpu(workspaceSize, tensorrt_llm::DataType::kINT32); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void PenaltyLayer::allocateBuffer() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mLogitsPtrsHost = mBufferManager->pinnedPool(ITensor::makeShape({}), TRTDataType::value); - auto const batchSizeShape = ITensor::makeShape({mDecoderDomain.getBatchSize()}); - mTemperature = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mRepetitionPenalty = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mPresencePenalty = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mFrequencyPenalty = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mMinLength = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mPromptIgnoreLength = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - - if (mDecodingMode.isUseTemperature()) - { - mTemperatureDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); - } - if (mDecodingMode.isUseRepetitionPenalty()) - { - mRepetitionPenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); - } - if (mDecodingMode.isUsePresencePenalty()) - { - mPresencePenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); - } - if (mDecodingMode.isUseFrequencyPenalty()) - { - mFrequencyPenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); - } - if (mDecodingMode.isUseMinLength()) - { - mMinLengthDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kINT32); - } - if (mDecodingMode.isUseOccurrencePenalty()) - { - mPromptIgnoreLengthDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kINT32); - } - - auto const logitsPtrDeviceDesc = std::make_pair(batchSizeShape, TRTDataType::value); - mWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize(logitsPtrDeviceDesc); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void PenaltyLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(PenaltyLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - if (mConfiguredBeamWidth == -1) - { - // This code is left only for Python runtime - // In C++ runtime given maxBeamWidth should always be equal to the runtime beamWidth - TLLM_CHECK(mDecodingMode.isAuto()); - mConfiguredBeamWidth = beamWidth; - mDecodingMode - = mConfiguredBeamWidth == 1 ? executor::DecodingMode::TopKTopP() : executor::DecodingMode::BeamSearch(); - allocateWorkspace(); - } - - // Setup penalties. - FillBuffers const fillBuffers{batchSize, mDecoderDomain.getBatchSize(), mBufferManager}; - - auto const& penaltyParams = setupParams->penaltyParams; - TLLM_CHECK_WITH_INFO(penaltyParams, "penaltyParams for setup is not set"); - - bool const useTemperature = mDecodingMode.isUseTemperature() && penaltyParams->temperature.has_value(); - bool const useRepetitionPenalty - = mDecodingMode.isUseRepetitionPenalty() && penaltyParams->repetitionPenalty.has_value(); - bool const usePresencePenalty = mDecodingMode.isUsePresencePenalty() && penaltyParams->presencePenalty.has_value(); - bool const useFrequencyPenalty - = mDecodingMode.isUseFrequencyPenalty() && penaltyParams->frequencyPenalty.has_value(); - bool const useMinLength = mDecodingMode.isUseMinLength() && penaltyParams->minLength.has_value(); - bool const usePromptIgnoreLength - = mDecodingMode.isUseOccurrencePenalty() && penaltyParams->promptIgnoreLength.has_value(); - // FIXME: once one of the requests has some penalty, we will always have to compute it. - // To avoid that we need to scan through all active requests at each iteration. - mUseTemperature |= useTemperature; - mUseRepetitionPenalty |= useRepetitionPenalty; - mUsePresencePenalty |= usePresencePenalty; - mUseFrequencyPenalty |= useFrequencyPenalty; - mUseMinLength |= useMinLength; - mUsePromptIgnoreLength |= usePromptIgnoreLength; - - if (mUseTemperature) - { - fillBuffers(penaltyParams->temperature, DefaultDecodingParams::getTemperature(), mTemperature, - mTemperatureDevice, batchSlots, getLimitsPenalty(DecodingPenaltyType::Temperature), "temperature penalty"); - } - if (mUseRepetitionPenalty) - { - fillBuffers(penaltyParams->repetitionPenalty, DefaultDecodingParams::getRepetitionPenalty(), mRepetitionPenalty, - mRepetitionPenaltyDevice, batchSlots, getLimitsPenalty(DecodingPenaltyType::Repetition), - "repetition penalty"); - } - if (mUsePresencePenalty) - { - fillBuffers(penaltyParams->presencePenalty, DefaultDecodingParams::getPresencePenalty(), mPresencePenalty, - mPresencePenaltyDevice, batchSlots, getLimitsPenalty(DecodingPenaltyType::Presence), "presence penalty"); - } - if (mUseFrequencyPenalty) - { - fillBuffers(penaltyParams->frequencyPenalty, DefaultDecodingParams::getFrequencyPenalty(), mFrequencyPenalty, - mFrequencyPenaltyDevice, batchSlots, getLimitsPenalty(DecodingPenaltyType::Frequency), "frequency penalty"); - } - if (mUseMinLength) - { - fillBuffers(penaltyParams->minLength, DefaultDecodingParams::getMinLength(), mMinLength, mMinLengthDevice, - batchSlots, getLimitsPenalty(DecodingPenaltyType::MinLength), "min length"); - } - if (mUsePromptIgnoreLength) - { - fillBuffers(penaltyParams->promptIgnoreLength, DefaultDecodingParams::getPromptIgnoreLength(), - mPromptIgnoreLength, mPromptIgnoreLengthDevice, batchSlots, - getLimitsPenalty(DecodingPenaltyType::PromptIgnoreLength), "prompt ignore length"); - } - - // Reset penalty workspace - auto const workspaceSizePerBatch - = mDecoderDomain.getMaxDecodingTokens() * mConfiguredBeamWidth * mDecoderDomain.getVocabSize() * 2; - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto batchSlot = runtime::bufferCast(*batchSlots)[bi]; - - if (mPenaltyWorkspaceDevice) - { - auto deviceSlice = runtime::IBuffer::slice( - mPenaltyWorkspaceDevice, batchSlot * workspaceSizePerBatch, workspaceSizePerBatch); - mBufferManager->setZero(*deviceSlice); - } - - if (mPenaltyWorkspacePrevDevice) - { - auto deviceSlice = runtime::IBuffer::slice( - mPenaltyWorkspacePrevDevice, batchSlot * workspaceSizePerBatch, workspaceSizePerBatch); - mBufferManager->setZero(*deviceSlice); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void PenaltyLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(PenaltyLayer_forwardAsync); - - auto outputs = std::dynamic_pointer_cast(baseOutputs); - auto params = std::dynamic_pointer_cast(baseInputs); - - auto const localDecoderDomain = getLocalDecoderDomain(params, mDecoderDomain); - auto const maxSeqLen = outputs->outputIds->getDimension<-1>(); - - if (mLogitsPtrsHost->data() == nullptr) - { - mLogitsPtrsHost->reshape( - ITensor::makeShape({static_cast(maxSeqLen), static_cast(mDecoderDomain.getBatchSize())})); - mRuntimeMaxSeqLen = maxSeqLen; - } - - mCyclicStep = mCyclicStep % mRuntimeMaxSeqLen; - - TensorPtr logitsPtrsHost = ITensor::slice(mLogitsPtrsHost, mCyclicStep, 1); - logitsPtrsHost->squeeze(0); - auto logitsPtrsHostData = bufferCast(*logitsPtrsHost); - for (SizeType32 bi = 0; bi < localDecoderDomain.getBatchSize(); bi++) - { - if (params->logitsVec) - { - TLLM_CHECK_WITH_INFO(params->logitsVec->size() == static_cast(localDecoderDomain.getBatchSize()), - "Logits vector size (%lu) is not equal to the batchSize (%d)", params->logitsVec->size(), - localDecoderDomain.getBatchSize()); - logitsPtrsHostData[bi] = bufferCastOrNull(params->logitsVec.value()[bi]); - } - else - { - TensorConstPtr logitsForBatchIndex = ITensor::slice(params->logits.value(), ITensor::makeShape({bi})); - auto const ptrToLogitsForBatchIndex = bufferCastOrNull(logitsForBatchIndex); - logitsPtrsHostData[bi] = ptrToLogitsForBatchIndex; - } - } - - auto const* inputLengths = bufferCastOrNull(params->inputLengths); - auto embeddingBias = bufferCastOrNull(params->embeddingBias); - auto const* batchSlotsHostPtr = bufferCast(*params->batchSlots); -#define GET_PENALTIES(capital_name, type) \ - (mUse##capital_name \ - && !allOfBatchSlots(batchSlotsHostPtr, bufferCast(*m##capital_name), localDecoderDomain.getBatchSize(), \ - DefaultDecodingParams::get##capital_name())) \ - ? m##capital_name##Device \ - : nullptr; - - auto temperatures = GET_PENALTIES(Temperature, float); - auto repetitionPenalties = GET_PENALTIES(RepetitionPenalty, float); - auto presencePenalties = GET_PENALTIES(PresencePenalty, float); - auto frequencyPenalties = GET_PENALTIES(FrequencyPenalty, float); - auto minLengths = GET_PENALTIES(MinLength, SizeType32); - auto promptIgnoreLengths = GET_PENALTIES(PromptIgnoreLength, SizeType32); - -#undef GET_PENALTIES - - auto* const tokensPerStep = bufferCastOrNull(params->curTokensPerStep); - - InvokeBatchApplyPenaltyParams penaltyParams{}; - - TensorPtr logitsPtrsHostSlice = ITensor::slice(logitsPtrsHost, 0, localDecoderDomain.getBatchSize()); - auto [logitsPtrsDeviceSlice] = workspace->mirrorInWorkspace(logitsPtrsHostSlice); - auto runtimeLogits = workspace->getDeviceRuntimeLogits(); - penaltyParams.inputLogits = reinterpret_cast(bufferCast(*logitsPtrsDeviceSlice)); - penaltyParams.outputLogits = bufferCast(*runtimeLogits); - penaltyParams.biases = embeddingBias; - penaltyParams.penaltyWorkspace = bufferCastOrNull(mPenaltyWorkspaceDevice); - penaltyParams.penaltyWorkspacePrev = bufferCastOrNull(mPenaltyWorkspacePrevDevice); - penaltyParams.temperatures = bufferCastOrNull(temperatures); - penaltyParams.repetitionPenalties = bufferCastOrNull(repetitionPenalties); - penaltyParams.presencePenalties = bufferCastOrNull(presencePenalties); - penaltyParams.frequencyPenalties = bufferCastOrNull(frequencyPenalties); - penaltyParams.batchSize = localDecoderDomain.getBatchSize(); - penaltyParams.beamWidth = localDecoderDomain.getBeamWidth(); - penaltyParams.maxSeqLen = maxSeqLen; - penaltyParams.vocabSize = mDecoderDomain.getVocabSize(); - penaltyParams.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - penaltyParams.outputIdsPtr = bufferCast(*outputs->outputIdsPtr); - penaltyParams.parentIdsPtr = bufferCast(*outputs->parentIdsPtr); - penaltyParams.inputLengths = inputLengths; - penaltyParams.sequenceLengths = bufferCast(*outputs->sequenceLength.value()); - penaltyParams.minLengths = bufferCastOrNull(minLengths); - penaltyParams.promptIgnoreLengths = bufferCastOrNull(promptIgnoreLengths); - penaltyParams.endIds = bufferCast(*params->endIds); - penaltyParams.batchSlots = workspace->getDeviceBatchSlotsPtr(); - penaltyParams.maxTokensPerStep = mDecoderDomain.getMaxDecodingTokens(); - penaltyParams.tokensPerStep = tokensPerStep; - penaltyParams.finished = (params->finished) - ? reinterpret_cast(bufferCast(*params->finished.value())) - : nullptr; - penaltyParams.stream = getStream(); - - if (penaltyParams.beamWidth > 1) - { - // Convert logits into logProbs before penalties, only necessary in Beam-Search. - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logitsPtrs = const_cast(penaltyParams.inputLogits); - biasSoftmaxParams.bias = penaltyParams.biases; - biasSoftmaxParams.endIds = penaltyParams.endIds; - biasSoftmaxParams.batchSlots = penaltyParams.batchSlots; - biasSoftmaxParams.batchSize = penaltyParams.batchSize; - biasSoftmaxParams.maxBatchSize = mDecoderDomain.getBatchSize(); - biasSoftmaxParams.maxBeamWidth = penaltyParams.beamWidth; - biasSoftmaxParams.vocabSize = penaltyParams.vocabSize; - biasSoftmaxParams.vocabSizePadded = penaltyParams.vocabSizePadded; - biasSoftmaxParams.skipSoftMax = false; - biasSoftmaxParams.batchSlotsLogits = penaltyParams.batchSlots != nullptr; - biasSoftmaxParams.checkParams(); - invokeAddBiasSoftMax(biasSoftmaxParams, penaltyParams.stream); - } - - invokeBatchApplyPenalty(penaltyParams); - sync_check_cuda_error(penaltyParams.stream); - - mCyclicStep += 1; - - auto const logitsShape = ITensor::makeShape({localDecoderDomain.getBatchSize(), - mDecoderDomain.getMaxDecodingTokens(), localDecoderDomain.getBeamWidth(), mDecoderDomain.getVocabSizePadded()}); - params->logits = ITensor::view(runtimeLogits, logitsShape); - - if (mDecodingMode.isBeamSearch()) - { - std::swap(mPenaltyWorkspaceDevice, mPenaltyWorkspacePrevDevice); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class PenaltyLayer; -template class PenaltyLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/penaltyLayer.h b/cpp/tensorrt_llm/layers/penaltyLayer.h deleted file mode 100644 index 7fa9e3a38952..000000000000 --- a/cpp/tensorrt_llm/layers/penaltyLayer.h +++ /dev/null @@ -1,95 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" - -namespace tensorrt_llm::layers -{ - -//! \brief Layer applies penalties to the logits. Supports: -//! 1. Temperature -//! 2. Repetition penalty -//! 3. Presence penalty -//! 4. Frequency penalty -//! 5. Min length penalty -template -class PenaltyLayer : public BaseLayer -{ -public: - PenaltyLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - //! \brief Modifies 'outputs->logits' in-place with -INF for banned words - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - void initialize(); - void allocateWorkspace(); - void allocateBuffer(); - -private: - using BaseLayer::mDecoderDomain; - - executor::DecodingMode mDecodingMode; - - size_t mWorkspaceSize{}; - TensorPtr mTemperatureDevice; - TensorPtr mRepetitionPenaltyDevice; - TensorPtr mPresencePenaltyDevice; - TensorPtr mFrequencyPenaltyDevice; - TensorPtr mMinLengthDevice; - TensorPtr mPromptIgnoreLengthDevice; - - TensorPtr mTemperature; - TensorPtr mRepetitionPenalty; - TensorPtr mPresencePenalty; - TensorPtr mFrequencyPenalty; - TensorPtr mMinLength; - TensorPtr mPromptIgnoreLength; - - bool mUseTemperature{false}; - bool mUseRepetitionPenalty{false}; - bool mUsePresencePenalty{false}; - bool mUseFrequencyPenalty{false}; - bool mUseMinLength{false}; - bool mUsePromptIgnoreLength{false}; - - runtime::SizeType32 mCyclicStep{0}; - runtime::SizeType32 mRuntimeMaxSeqLen{0}; - runtime::SizeType32 mConfiguredBeamWidth{-1}; - - BufferPtr mPenaltyWorkspaceDevice; - BufferPtr mPenaltyWorkspacePrevDevice; - TensorPtr mLogitsPtrsHost; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/samplingLayer.cpp b/cpp/tensorrt_llm/layers/samplingLayer.cpp deleted file mode 100644 index fe8166f986e1..000000000000 --- a/cpp/tensorrt_llm/layers/samplingLayer.cpp +++ /dev/null @@ -1,213 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" -#include "tensorrt_llm/layers/topKSamplingLayer.h" -#include "tensorrt_llm/layers/topPSamplingLayer.h" - -#include "samplingLayer.h" -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -SamplingLayer::SamplingLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO(!mDecodingMode.isBeamSearch(), "SamplingLayer does not support Beam search mode"); - TLLM_CHECK_WITH_INFO(mDecodingMode.isTopKorTopP(), "SamplingLayer requires TopK or TopP mode"); - if (mDecodingMode.isTopK()) - { - mSamplingLayers.emplace_back(std::make_unique>(decoderDomain, mBufferManager)); - } - - if (mDecodingMode.isTopP()) - { - mSamplingLayers.emplace_back( - std::make_unique>(decoderDomain, mBufferManager, /* deterministic */ true)); - } - - allocateBuffer(decoderDomain.getBatchSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void SamplingLayer::allocateBuffer(SizeType32 batchSize) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - size_t workspaceSize = 0; - for (auto&& layer : mSamplingLayers) - { - workspaceSize = std::max(workspaceSize, layer->getWorkspaceSize()); - } - mWorkspaceSize = workspaceSize; - - auto const batchSizeShape = ITensor::makeShape({batchSize}); - mSetupWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(batchSizeShape, TRTDataType::value)); - mSkipDecodeDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({batchSize, sizeof(curandState_t)}), TRTDataType::value); - - // host buffers. - mSkipDecodeHost = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - - mRuntimeMinPHost = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - mRuntimeMinPDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - TLLM_CHECK(mSkipDecodeHost != nullptr); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void SamplingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(SamplingLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - workspace->initializeDeviceCurandStates( - setupParams->randomSeed, batchSize, workspace->getDeviceBatchSlots(), mCurandStatesDevice); - - if (setupParams->outputLogProbs) - { - // FIXME: monotonically growing - mOutputLogProbs = std::any_of(setupParams->outputLogProbs->begin(), setupParams->outputLogProbs->end(), - [this](bool outputLogProbs) { return this->mOutputLogProbs | outputLogProbs; }); - } - - if (setupParams->cumLogProbs) - { - // FIXME: monotonically growing - mCumLogProbs = std::any_of(setupParams->cumLogProbs->begin(), setupParams->cumLogProbs->end(), - [this](bool cumLogProbs) { return this->mCumLogProbs | cumLogProbs; }); - } - - for (auto&& layer : mSamplingLayers) - { - layer->setup(batchSize, beamWidth, batchSlots, setupParams, workspace); - } - - FillBuffers const fillBuffers{batchSize, mDecoderDomain.getBatchSize(), mBufferManager}; - bool const useMinP = mDecodingMode.isUseMinP() && setupParams->runtimeMinP.has_value(); - mUseMinP |= useMinP; - if (mUseMinP) - { - fillBuffers(setupParams->runtimeMinP, DefaultDecodingParams::getMinP(), mRuntimeMinPHost, mRuntimeMinPDevice, - batchSlots, std::pair(-1e-6f, 1.0f), "min_p"); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void SamplingLayer::forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(SamplingLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto const localDecoderDomain = getLocalDecoderDomain(inputs, mDecoderDomain); - auto const batchSize = inputs->logits.value()->getDimension<0>(); - - auto const* endIds = bufferCast(*inputs->endIds); - - FinishedState const* finishedInput = (inputs->finished) - ? reinterpret_cast(bufferCast(*inputs->finished.value())) - : nullptr; - - auto const skipTopP = !mDecodingMode.isTopP(); - - auto const* batchSlotsHostPtr = bufferCast(*inputs->batchSlots); - auto minPs = mUseMinP - && !allOfBatchSlots(batchSlotsHostPtr, bufferCast(*mRuntimeMinPHost), - localDecoderDomain.getBatchSize(), DefaultDecodingParams::getMinP()) - ? mRuntimeMinPDevice - : nullptr; - - // Compute probabilities either for TopP or if cumLogProbs or outputLogProbs are specified - bool const skipSoftMax = skipTopP && !mOutputLogProbs && !mCumLogProbs && minPs == nullptr; - - inputs->curandStates = reinterpret_cast(bufferCast(*mCurandStatesDevice)); - inputs->probsComputed = !skipSoftMax; - if (!skipSoftMax) - { - auto runtimeLogitsPtr = bufferCast(*workspace->getDeviceRuntimeLogits()); - auto logitsPtrsPtr = static_cast(nullptr); - auto biasPtr = static_cast(nullptr); - auto const* batchSlotsPtr = workspace->getDeviceBatchSlotsPtr(); - - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logits = runtimeLogitsPtr; - biasSoftmaxParams.logitsPtrs = logitsPtrsPtr; - biasSoftmaxParams.probs = runtimeLogitsPtr; - biasSoftmaxParams.bias = biasPtr; - biasSoftmaxParams.endIds = endIds; - biasSoftmaxParams.finished = finishedInput; - biasSoftmaxParams.batchSlots = batchSlotsPtr; - biasSoftmaxParams.batchSize = batchSize; - biasSoftmaxParams.maxBatchSize = mDecoderDomain.getBatchSize(); - biasSoftmaxParams.maxBeamWidth = 1; - biasSoftmaxParams.vocabSize = mDecoderDomain.getVocabSize(); - biasSoftmaxParams.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - biasSoftmaxParams.skipSoftMax = skipSoftMax; - biasSoftmaxParams.batchSlotsLogits = false; - biasSoftmaxParams.minPs = bufferCastOrNull(minPs); - biasSoftmaxParams.checkParams(); - invokeAddBiasSoftMax(biasSoftmaxParams, getStream()); - sync_check_cuda_error(getStream()); - } - - for (auto&& layer : mSamplingLayers) - { - layer->forwardAsync(outputs, baseInputs, workspace); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t SamplingLayer::getWorkspaceSize() const noexcept -{ - return std::max(mWorkspaceSize, mSetupWorkspaceSize); -} - -template class SamplingLayer; -template class SamplingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/samplingLayer.h b/cpp/tensorrt_llm/layers/samplingLayer.h deleted file mode 100644 index 6fc17f495c04..000000000000 --- a/cpp/tensorrt_llm/layers/samplingLayer.h +++ /dev/null @@ -1,79 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/common.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Top class for sampling layers. -//! It sets up and executes TopKSamplingLayer and TopPSamplingLayer samplings -template -class SamplingLayer : public BaseLayer -{ -public: - using Base = BaseLayer; - - SamplingLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - using Base::mDecoderDomain; - - executor::DecodingMode mDecodingMode; - - size_t mWorkspaceSize{0}; - size_t mSetupWorkspaceSize{0}; - - TensorPtr mCurandStatesDevice; - TensorPtr mSkipDecodeDevice; - - TensorPtr mSkipDecodeHost; - bool mSkipAny{false}; - - bool mOutputLogProbs{false}; - bool mCumLogProbs{false}; - - TensorPtr mRuntimeMinPHost; - TensorPtr mRuntimeMinPDevice; - bool mUseMinP{false}; - - std::vector> mSamplingLayers; - -private: - void allocateBuffer(runtime::SizeType32 batchSize); -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/stopCriteriaLayer.cpp b/cpp/tensorrt_llm/layers/stopCriteriaLayer.cpp deleted file mode 100644 index fbe7a2fb2f11..000000000000 --- a/cpp/tensorrt_llm/layers/stopCriteriaLayer.cpp +++ /dev/null @@ -1,179 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "stopCriteriaLayer.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/stopCriteriaKernels.h" -#include "tensorrt_llm/layers/layerUtils.h" - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -size_t StopCriteriaLayer::getWorkspaceSize() const noexcept -{ - return mWorkspaceSize; -} - -template -StopCriteriaLayer::StopCriteriaLayer(executor::DecodingMode const& mode, DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) - , mDecodingMode(mode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const stopWordsWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(ITensor::makeShape({decoderDomain.getBatchSize()}), TRTDataType::value), - std::make_pair(ITensor::makeShape({decoderDomain.getBatchSize()}), TRTDataType::value), - std::make_pair(ITensor::makeShape({decoderDomain.getBatchSize(), decoderDomain.getBeamWidth()}), - TRTDataType::value)); - auto const lengthCriterionWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(ITensor::makeShape({1}), TRTDataType::value), - std::make_pair(ITensor::makeShape({decoderDomain.getBatchSize(), decoderDomain.getBeamWidth()}), - TRTDataType::value)); - mWorkspaceSize = std::max(stopWordsWorkspaceSize, lengthCriterionWorkspaceSize); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void StopCriteriaLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void StopCriteriaLayer::forwardAsync(std::shared_ptr const& baseOutputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(StopCriteriaLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - auto outputs = std::dynamic_pointer_cast(baseOutputs); - - auto localDecoderDomain = getLocalDecoderDomain(inputs, mDecoderDomain); - // Beam width might have been changed in Variable-Beam-Width-Search mode - localDecoderDomain.setBeamWidth(baseOutputs->beamWidth); - - auto const maxSeqLen = outputs->outputIds->getDimension<-1>(); - - TLLM_CHECK_WITH_INFO(inputs->stopCriteriaInputs, "stopCriteriaInputs for forward is not set"); - - if (mDecodingMode.isUseStopWords() && inputs->stopCriteriaInputs->maxStopWordsLen != 0) - { - checkStopWordsStopCriteria(outputs, inputs, localDecoderDomain, maxSeqLen, *mBufferManager, workspace); - } - if (mDecodingMode.isUseExplicitEosStop()) - { - checkEosToken(outputs, inputs, localDecoderDomain, *mBufferManager, workspace); - } - if (mDecodingMode.isUseMaxLengthStop() && inputs->stopCriteriaInputs->sequenceLimitLength) - { - checkMaxLengthStopCriteria(outputs, inputs, localDecoderDomain, *mBufferManager, workspace); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void StopCriteriaLayer::checkStopWordsStopCriteria(std::shared_ptr& outputs, - std::shared_ptr const& inputs, DecoderDomain const& decoderDomain, SizeType32 maxSeqLen, - BufferManager const& bufferManager, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const maxStopWordsLength = inputs->stopCriteriaInputs->maxStopWordsLen; - auto* numNewTokens = bufferCastOrNull(outputs->numNewTokens); - auto* outputIdsPtr = bufferCast(*outputs->outputIdsPtr); - auto* parentIdsPtr = bufferCast(*outputs->parentIdsPtr); - auto* sequenceLengthPtr = bufferCastOrNull(outputs->sequenceLength); - auto [stopWordsLengthsDevice, stopWordsPtrDevice, finishedDevice] - = workspace->mirrorInWorkspace(inputs->stopCriteriaInputs->stopWordsLengths.value_or(nullptr), - inputs->stopCriteriaInputs->stopWordsPtr.value_or(nullptr), outputs->finished.value_or(nullptr)); - auto const* stopWordsLengthsPtr - = stopWordsLengthsDevice == nullptr ? nullptr : bufferCast(*stopWordsLengthsDevice); - auto const* stopWordsPtrPtr - = stopWordsPtrDevice == nullptr ? nullptr : bufferCast(*stopWordsPtrDevice); - auto* finishedPtr = finishedDevice == nullptr - ? nullptr - : reinterpret_cast(bufferCast(*finishedDevice)); - invokeStopWordsCriterion(outputIdsPtr, parentIdsPtr, stopWordsPtrPtr, finishedPtr, sequenceLengthPtr, - workspace->getDeviceBatchSlotsPtr(), stopWordsLengthsPtr, numNewTokens, maxStopWordsLength, - decoderDomain.getBatchSize(), decoderDomain.getBeamWidth(), maxSeqLen, bufferManager.getStream().get()); - if (finishedPtr != nullptr) - { - bufferManager.copy(*finishedDevice, *outputs->finished.value()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void StopCriteriaLayer::checkMaxLengthStopCriteria(std::shared_ptr& outputs, - std::shared_ptr const& inputs, DecoderDomain const& decoderDomain, - BufferManager const& bufferManager, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto* numNewTokens = bufferCastOrNull(outputs->numNewTokens); - auto [finishedSumDevice, finishedDevice] - = workspace->mirrorInWorkspace(outputs->finishedSum.value_or(nullptr), outputs->finished.value_or(nullptr)); - auto* finishedSumDevicePtr = finishedSumDevice == nullptr ? nullptr : bufferCast(*finishedSumDevice); - auto* finishedPtr = finishedDevice == nullptr - ? nullptr - : reinterpret_cast(bufferCast(*finishedDevice)); - invokeLengthCriterion(finishedPtr, finishedSumDevicePtr, - bufferCastOrNull(inputs->stopCriteriaInputs->sequenceLimitLength), - bufferCastOrNull(outputs->sequenceLength), numNewTokens, workspace->getDeviceBatchSlotsPtr(), - decoderDomain.getBatchSize(), decoderDomain.getBeamWidth(), bufferManager.getStream().get()); - if (finishedSumDevice != nullptr) - { - bufferManager.copy(*finishedSumDevice, *outputs->finishedSum.value()); - } - if (finishedPtr != nullptr) - { - bufferManager.copy(*finishedDevice, *outputs->finished.value()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void StopCriteriaLayer::checkEosToken(std::shared_ptr& outputs, - std::shared_ptr const& inputs, DecoderDomain const& decoderDomain, - BufferManager const& bufferManager, std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto* numNewTokens = bufferCastOrNull(outputs->numNewTokens); - auto* sequenceLengthsPtr = bufferCastOrNull(outputs->sequenceLength); - auto const* endIdsPtr = bufferCastOrNull(inputs->endIds); - auto* finishedStatePtr - = reinterpret_cast(bufferCastOrNull(outputs->finished)); - invokeExplicitEOSCriterion(bufferCastOrNull(outputs->outputIdsPtr), endIdsPtr, finishedStatePtr, - sequenceLengthsPtr, numNewTokens, workspace->getDeviceBatchSlotsPtr(), decoderDomain.getBatchSize(), - decoderDomain.getBeamWidth(), decoderDomain.getMaxDecodingTokens(), bufferManager.getStream().get()); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template class StopCriteriaLayer; -template class StopCriteriaLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/stopCriteriaLayer.h b/cpp/tensorrt_llm/layers/stopCriteriaLayer.h deleted file mode 100644 index 4525621bf042..000000000000 --- a/cpp/tensorrt_llm/layers/stopCriteriaLayer.h +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/layers/decodingParams.h" - -#include - -namespace tensorrt_llm::layers -{ - -//! \brief Layer to process stop criteria. Supports: -//! 1. Stop words criteria -//! 2. Maximum length criteria -template -class StopCriteriaLayer : public BaseLayer -{ -public: - StopCriteriaLayer(executor::DecodingMode const& mode, DecoderDomain const& /* decoderDomain */, - std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -private: - static void checkMaxLengthStopCriteria(std::shared_ptr& outputs, - std::shared_ptr const& inputs, DecoderDomain const& decoderDomain, - runtime::BufferManager const& bufferManager, std::shared_ptr const& workspace); - static void checkStopWordsStopCriteria(std::shared_ptr& outputs, - std::shared_ptr const& inputs, DecoderDomain const& decoderDomain, - runtime::SizeType32 maxSeqLen, runtime::BufferManager const& bufferManager, - std::shared_ptr const& workspace); - static void checkEosToken(std::shared_ptr& outputs, - std::shared_ptr const& inputs, DecoderDomain const& decoderDomain, - runtime::BufferManager const& bufferManager, std::shared_ptr const& workspace); - -private: - using BaseLayer::mDecoderDomain; - - executor::DecodingMode mDecodingMode; - size_t mWorkspaceSize{0}; -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/topKSamplingLayer.cpp b/cpp/tensorrt_llm/layers/topKSamplingLayer.cpp deleted file mode 100644 index 1f4a62f3687a..000000000000 --- a/cpp/tensorrt_llm/layers/topKSamplingLayer.cpp +++ /dev/null @@ -1,200 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "topKSamplingLayer.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" - -#include -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -TopKSamplingLayer::TopKSamplingLayer( - DecoderDomain const& decoderDomain, std::shared_ptr bufferManager) - : BaseLayer(decoderDomain, bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - allocateBuffer(mDecoderDomain.getBatchSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void TopKSamplingLayer::allocateBuffer(SizeType32 const batchSize) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mWorkspaceSize = getTopKWorkspaceSize(batchSize, 1, TOP_K_MAX, mDecoderDomain.getVocabSizePadded()); - auto const batchSizeShape = ITensor::makeShape({batchSize}); - mRuntimeTopKDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mRuntimeTopPDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mSkipDecodeDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mRuntimeTopKHost = mBufferManager->cpu(batchSizeShape, TRTDataType::value); - mSkipDecodeHost = mBufferManager->cpu(batchSizeShape, TRTDataType::value); - mSetupWorkspaceSize = batchSize * sizeof(SizeType32); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void TopKSamplingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(TopKSamplingLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - mNormalizeLogProbs = setupParams->normalizeLogProbs.value_or(false); - - auto runtimeTopK = setupParams->runtimeTopK.value_or(std::vector{DefaultDecodingParams::getTopK()}); - auto runtimeTopP = setupParams->runtimeTopP.value_or(std::vector{DefaultDecodingParams::getTopP()}); - - auto const paramsSize = expandMatchElements(batchSize, runtimeTopK, runtimeTopP); - - TLLM_CHECK_WITH_INFO(paramsSize != 0, - fmtstr("TopKSamplingLayer got parameter with unexpected size, want 1 or batchSize(%d), got" - "runtimeTopK.size() = %zu, runtimeTopP.size() = %zu", - batchSize, runtimeTopK.size(), runtimeTopP.size())); - - for (size_t i = 0; i < paramsSize; ++i) - { - auto& topK = runtimeTopK[i]; - auto& topP = runtimeTopP[i]; - clampTopK(topK); - clampTopP(topP); - regularizeTopKTopP(topK, topP); - } - - // Update parameters on both device and host, so we can - // - determine whether we can skip launch kernel by examine mSkipDecodeHost - // - select best kernel by examine mRuntimeTopKHost - // without consulting device memory, or we'll have to do an expensive synchronization. - SizeType32* topKsPtr = nullptr; - float* topPsPtr = nullptr; - - if (paramsSize > 1) - { - auto initWorkspaceSizes = getTopKInitWorkspaceSizes(batchSize); - auto workspacePtr = workspace->getRawWorkspaceDevicePtr(); - calcAlignedPointers(workspacePtr, initWorkspaceSizes)(topKsPtr, topPsPtr); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, runtimeTopK, IBuffer::wrap(topKsPtr, initWorkspaceSizes[0] / sizeof(*topKsPtr))); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, runtimeTopP, IBuffer::wrap(topPsPtr, initWorkspaceSizes[1] / sizeof(*topPsPtr))); - } - auto const* batchSlotsDevicePtr = workspace->getDeviceBatchSlotsPtr(); - auto* skipDecodeDevicePtr = bufferCastOrNull(mSkipDecodeDevice); - invokeSetupTopKRuntimeArgs(batchSize, // - {topKsPtr, runtimeTopK.front(), bufferCast(*mRuntimeTopKDevice)}, // - {topPsPtr, runtimeTopP.front(), bufferCast(*mRuntimeTopPDevice)}, // - skipDecodeDevicePtr, batchSlotsDevicePtr, true, getStream()); - - auto const* batchSlotsHostPtr = bufferCast(*batchSlots); - auto* skipDecodeHostPtr = bufferCastOrNull(mSkipDecodeHost); - topKsPtr = paramsSize > 1 ? runtimeTopK.data() : nullptr; - invokeSetupTopKRuntimeArgs(batchSize, // - {topKsPtr, runtimeTopK.front(), bufferCast(*mRuntimeTopKHost)}, {}, // - skipDecodeHostPtr, batchSlotsHostPtr, false); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void TopKSamplingLayer::forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(TopKSamplingLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto const batchSize = inputs->logits.value()->getDimension<0>(); - - auto const* batchSlotsHost = bufferCast(*inputs->batchSlots); - auto* skipDecodeHostPtr = bufferCastOrNull(mSkipDecodeHost); - auto const skip = allOfBatchSlots(batchSlotsHost, skipDecodeHostPtr, batchSize, true); - if (skip) - { - return; - } - - auto logits = bufferCastOrNull(inputs->logits); - auto const* endIds = bufferCastOrNull(inputs->endIds); - auto const probsComputed = inputs->probsComputed; - - FinishedState const* finishedInput = (inputs->finished) - ? reinterpret_cast(bufferCastOrNull(inputs->finished)) - : nullptr; - FinishedState* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished)) - : nullptr; - - auto* runtimeTopKHostPtr = bufferCast(*mRuntimeTopKHost); - - TopKSamplingKernelParams params; - params.logProbs = logits; - params.outputIdsPtrs = bufferCastOrNull(outputs->outputIdsPtr); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.maxTopP = 1.0f; - params.topPs = bufferCastOrNull(mRuntimeTopPDevice); - params.maxTopK = maxOfBatchSlots(batchSlotsHost, runtimeTopKHostPtr, batchSize); - params.topKs = bufferCastOrNull(mRuntimeTopKDevice); - params.sequenceLengths = bufferCastOrNull(outputs->sequenceLength); - params.endIds = endIds; - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.finishedInput = finishedInput; - params.finishedOutput = finishedOutput; - params.skipDecode = bufferCastOrNull(mSkipDecodeDevice); - params.cumLogProbs = bufferCastOrNull(outputs->cumLogProbs); - params.outputLogProbs = bufferCastOrNull(outputs->outputLogProbsTiled); - params.curandState = inputs->curandStates; - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.maxTokensPerStep = 1; - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - params.normalizeLogProbs = mNormalizeLogProbs; - params.logitsHasProbs = probsComputed; - - invokeBatchTopKSampling(params, getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t TopKSamplingLayer::getWorkspaceSize() const noexcept -{ - return std::max(mWorkspaceSize, mSetupWorkspaceSize); -} - -template class TopKSamplingLayer; -template class TopKSamplingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/topKSamplingLayer.h b/cpp/tensorrt_llm/layers/topKSamplingLayer.h deleted file mode 100644 index e6a14f93bb1f..000000000000 --- a/cpp/tensorrt_llm/layers/topKSamplingLayer.h +++ /dev/null @@ -1,63 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/runtime/common.h" - -namespace tensorrt_llm::layers -{ - -//! \brief Layer to randomly sample tokens from TopK logits. -//! When both TopK and TopP are specified, layer jointly samples using TopK and TopP. -//! When no TopK param is specified, sampling is skipped for particular request. -template -class TopKSamplingLayer : public BaseLayer -{ - using Base = BaseLayer; - -public: - TopKSamplingLayer(DecoderDomain const& decoderDomain, std::shared_ptr bufferManager); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -protected: - bool mNormalizeLogProbs{true}; - size_t mWorkspaceSize{0}; - size_t mSetupWorkspaceSize{0}; - TensorPtr mRuntimeTopKDevice; - TensorPtr mRuntimeTopPDevice; - TensorPtr mSkipDecodeDevice; - TensorPtr mRuntimeTopKHost; - TensorPtr mSkipDecodeHost; - - using Base::mDecoderDomain; - -private: - void allocateBuffer(runtime::SizeType32 batchSize); -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/topPSamplingLayer.cpp b/cpp/tensorrt_llm/layers/topPSamplingLayer.cpp deleted file mode 100644 index eb826e4e5b04..000000000000 --- a/cpp/tensorrt_llm/layers/topPSamplingLayer.cpp +++ /dev/null @@ -1,311 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "topPSamplingLayer.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/samplingTopPKernels.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/layers/layerUtils.h" - -#include -#include - -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::layers -{ - -template -TopPSamplingLayer::TopPSamplingLayer(DecoderDomain const& decoderDomain, - std::shared_ptr bufferManager, bool isDeterministic, bool isAirTopP) - : BaseLayer(decoderDomain, bufferManager) - , mIsDeterministic(isDeterministic) - , mIsAirTopP(isAirTopP) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const deviceId = getDevice(); - TLLM_CUDA_CHECK(cudaGetDeviceProperties(&mDeviceProp, deviceId)); - - allocateBuffer(mDecoderDomain.getBatchSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void TopPSamplingLayer::allocateBuffer(SizeType32 batchSize) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (!mIsAirTopP) - { - mWorkspaceSize = getTopPWorkspaceSize(batchSize, mDecoderDomain.getVocabSizePadded()); - } - else - { - mWorkspaceSize = getAirTopPWorkspaceSize(batchSize, mDecoderDomain.getVocabSizePadded(), mIsDeterministic); - } - - auto const batchSizeShape = ITensor::makeShape({batchSize}); - mRuntimeTopKDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mRuntimeTopPDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mInitialTopPDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mTopPDecayDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mTopPMinDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mTopPResetIdsDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mSkipDecodeDevice = mBufferManager->gpu(batchSizeShape, TRTDataType::value); - mSkipDecodeHost = mBufferManager->pinnedPool(batchSizeShape, TRTDataType::value); - auto skipDecodeHostRange = BufferRange(*mSkipDecodeHost); - std::fill(skipDecodeHostRange.begin(), skipDecodeHostRange.end(), true); - - mSetupWorkspaceSize = std::max({mRuntimeTopKDevice->getSizeInBytes(), mRuntimeTopPDevice->getSizeInBytes(), - mInitialTopPDevice->getSizeInBytes(), mTopPDecayDevice->getSizeInBytes(), mTopPMinDevice->getSizeInBytes(), - mTopPResetIdsDevice->getSizeInBytes(), mSkipDecodeDevice->getSizeInBytes()}); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void TopPSamplingLayer::setup(SizeType32 batchSize, SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& baseSetupParams, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(TopPSamplingLayer_setup); - - auto setupParams = std::dynamic_pointer_cast(baseSetupParams); - - auto constexpr defaultTopPDecay = DefaultDecodingParams::getTopPDecay(); - auto constexpr defaultTopPMin = DefaultDecodingParams::getTopPMin(); // prevent TopP becoming 0.0 - - auto const* batchSlotsHostPtr = bufferCastOrNull(batchSlots); - auto* skipDecodeHostPtr = bufferCastOrNull(mSkipDecodeHost); - if (!setupParams->runtimeTopP.has_value() || setupParams->runtimeTopP.value().empty()) - { - // Fast path to disable TopP for slots - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto const bid = batchSlotsHostPtr[bi]; - skipDecodeHostPtr[bid] = true; - } - auto const maxBatchSize = mDecoderDomain.getBatchSize(); - auto skipDecodeHostSlice = IBuffer::slice(mSkipDecodeHost, 0, maxBatchSize); - mBufferManager->copy(*skipDecodeHostSlice, *mSkipDecodeDevice); - return; - } - - auto runtimeTopK = setupParams->runtimeTopK.value_or(std::vector{DefaultDecodingParams::getTopK()}); - auto runtimeTopP = setupParams->runtimeTopP.value(); - auto decayVec = setupParams->topPDecay.value_or(std::vector{defaultTopPDecay}); - auto topPMinVec = setupParams->topPMin.value_or(std::vector{defaultTopPMin}); - auto topPResetIdsVec = setupParams->topPResetIds.value_or(std::vector{DefaultDecodingParams::getTopPResetId()}); - - auto const paramsSize - = expandMatchElements(batchSize, runtimeTopK, runtimeTopP, decayVec, topPMinVec, topPResetIdsVec); - TLLM_CHECK_WITH_INFO(paramsSize != 0, - fmtstr("TopPSamplingLayer got parameter with unexpected size, want 1 or batchSize(%d), got" - "runtimeTopK.size() = %zu, " - "runtimeTopP.size() = %zu, " - "topPDecay.size() = %zu, " - "topPMin.size() = %zu, " - "topPResetIds.size() = %zu", - batchSize, runtimeTopK.size(), runtimeTopP.size(), decayVec.size(), topPMinVec.size(), - topPResetIdsVec.size())); - - for (size_t i = 0; i < paramsSize; ++i) - { - // support topK up to TOP_K_MAX. - auto& topK = runtimeTopK[i]; - auto& topP = runtimeTopP[i]; - clampTopK(topK); - clampTopP(topP); - regularizeTopKTopP(topK, topP); - - auto& decay = decayVec[i]; - if (decay <= 0.f || decay > 1.0f) - { - TLLM_LOG_WARNING( - "Decay (%f) is out of range ((0.0, 1.0f]). Change to default (%f).", decay, defaultTopPDecay); - decay = defaultTopPDecay; - } - - auto& topPMin = topPMinVec[i]; - if (topPMin <= 0.f || topPMin > 1.0f) - { - TLLM_LOG_WARNING( - "TopP min (%f) is out of range ([0.0, 1.0f]). Change to default (%f).", topPMin, defaultTopPMin); - topPMin = defaultTopPMin; - } - } - - // Update parameters on both device and host, so we can - // determine whether we can skip launch kernel by examine mSkipDecodeHost - // without consulting device memory, or we'll have to do an expensive synchronization. - SizeType32* topKsPtr = nullptr; - float* topPsPtr = nullptr; - float* topPDecayPtr = nullptr; - float* topPMinPtr = nullptr; - SizeType32* topPResetIdsPtr = nullptr; - - if (paramsSize > 1) - { - auto initWorkspaceSizes = getTopPInitWorkspaceSizes(batchSize); - std::vector alignedPointers; - calcAlignedPointers(workspace->getRawWorkspaceDevicePtr(), initWorkspaceSizes)( - topKsPtr, topPsPtr, topPDecayPtr, topPMinPtr, topPResetIdsPtr); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, runtimeTopK, IBuffer::wrap(topKsPtr, initWorkspaceSizes[0] / sizeof(*topKsPtr))); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, runtimeTopP, IBuffer::wrap(topPsPtr, initWorkspaceSizes[1] / sizeof(*topPsPtr))); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, decayVec, IBuffer::wrap(topPDecayPtr, initWorkspaceSizes[2] / sizeof(*topPDecayPtr))); - DecodingLayerWorkspace::copyToWorkspace( - *mBufferManager, topPMinVec, IBuffer::wrap(topPMinPtr, initWorkspaceSizes[3] / sizeof(*topPMinPtr))); - DecodingLayerWorkspace::copyToWorkspace(*mBufferManager, topPResetIdsVec, - IBuffer::wrap(topPResetIdsPtr, initWorkspaceSizes[4] / sizeof(*topPResetIdsPtr))); - } - - auto const* batchSlotsDevicePtr = workspace->getDeviceBatchSlotsPtr(); - auto* skipDecodeDevicePtr = bufferCastOrNull(mSkipDecodeDevice); - auto* initialTopPDevicePtr = bufferCast(*mInitialTopPDevice); - invokeSetTopPRuntimeArgs(batchSize, // - {topKsPtr, runtimeTopK.front(), bufferCast(*mRuntimeTopKDevice)}, // - {topPsPtr, runtimeTopP.front(), bufferCast(*mRuntimeTopPDevice)}, // - skipDecodeDevicePtr, initialTopPDevicePtr, batchSlotsDevicePtr, true, getStream()); - - invokeScatterDecodingParams(topPDecayPtr, decayVec.front(), bufferCast(*mTopPDecayDevice), - batchSlotsDevicePtr, batchSize, getStream()); - invokeScatterDecodingParams(topPMinPtr, topPMinVec.front(), bufferCast(*mTopPMinDevice), batchSlotsDevicePtr, - batchSize, getStream()); - invokeScatterDecodingParams(topPResetIdsPtr, topPResetIdsVec.front(), bufferCast(*mTopPResetIdsDevice), - batchSlotsDevicePtr, batchSize, getStream()); - - topKsPtr = paramsSize > 1 ? runtimeTopK.data() : nullptr; - invokeSetTopPRuntimeArgs(batchSize, // - {topKsPtr, runtimeTopK.front(), nullptr}, {}, // - skipDecodeHostPtr, nullptr, batchSlotsHostPtr, false); - - if (mIsAirTopP) - { - auto smCnt = mDeviceProp.multiProcessorCount; - if (smCnt <= 0) - { - auto const deviceId = getDevice(); - cudaDeviceProp prop{}; - TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, deviceId)); - smCnt = prop.multiProcessorCount; - } - mAirTopPBlockNum - = calcAirTopPBlockNum(batchSize, mDecoderDomain.getVocabSizePadded(), smCnt, mIsDeterministic); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void TopPSamplingLayer::forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& baseInputs, - std::shared_ptr const& workspace) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(TopPSamplingLayer_forwardAsync); - - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto const batchSize = inputs->logits.value()->getDimension<0>(); - - auto const* batchSlotsHost = bufferCast(*inputs->batchSlots); - auto* skipDecodeHostPtr = bufferCastOrNull(mSkipDecodeHost); - auto const skip = allOfBatchSlots(batchSlotsHost, skipDecodeHostPtr, batchSize, true); - if (skip) - { - return; - } - - // Probabilities must be already computed instead of logits - auto probs = bufferCastOrNull(inputs->logits); - auto const* endIds = bufferCastOrNull(inputs->endIds); - - auto const* finishedInput = (inputs->finished) ? reinterpret_cast( - bufferCastOrNull(inputs->finished.value())) - : nullptr; - auto* finishedOutput = (outputs->finished) - ? reinterpret_cast(bufferCastOrNull(outputs->finished.value())) - : nullptr; - - auto* cumLogProbs = bufferCastOrNull(outputs->cumLogProbs); - auto* outputLogProbs = bufferCastOrNull(outputs->outputLogProbsTiled); - auto* sequenceLength = bufferCastOrNull(outputs->sequenceLength); - - TopPSamplingKernelParams params{}; - params.probs = probs; - params.outputIdsPtrs = bufferCastOrNull(outputs->outputIdsPtr); - params.workspace = workspace->getRawWorkspaceDevicePtr(); - params.topPs = bufferCastOrNull(mRuntimeTopPDevice); - params.sequenceLength = sequenceLength; - params.endIds = endIds; - params.batchSlots = workspace->getDeviceBatchSlotsPtr(); - params.finishedInput = finishedInput; - params.finishedOutput = finishedOutput; - params.skipDecode = bufferCastOrNull(mSkipDecodeDevice); - params.cumLogProbs = cumLogProbs; - params.outputLogProbs = outputLogProbs; - params.curandState = inputs->curandStates; - params.batchSize = batchSize; - params.maxBatchSize = mDecoderDomain.getBatchSize(); - params.vocabSizePadded = mDecoderDomain.getVocabSizePadded(); - - if (!mIsAirTopP) - { - invokeBatchTopPSampling(params, getStream()); - } - else - { - params.blockNum = mAirTopPBlockNum; - params.isDeterministic = mIsDeterministic; - invokeBatchAirTopPSampling(params, getStream()); - } - - sync_check_cuda_error(getStream()); - auto* runtimeTopPDevicePtr = bufferCastOrNull(mRuntimeTopPDevice); - auto* initialTopPDevicePtr = bufferCastOrNull(mInitialTopPDevice); - auto* topPDecayDevicePtr = bufferCastOrNull(mTopPDecayDevice); - auto* topPMinDevicePtr = bufferCastOrNull(mTopPMinDevice); - auto* topPResetIdsDevice = bufferCastOrNull(mTopPResetIdsDevice); - auto* outputIdsPtr = bufferCastOrNull(outputs->outputIdsPtr); - invokeComputeToppDecay(runtimeTopPDevicePtr, initialTopPDevicePtr, outputIdsPtr, topPDecayDevicePtr, - topPMinDevicePtr, topPResetIdsDevice, sequenceLength, workspace->getDeviceBatchSlotsPtr(), batchSize, - getStream()); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -size_t TopPSamplingLayer::getWorkspaceSize() const noexcept -{ - return std::max(mSetupWorkspaceSize, mWorkspaceSize); -} - -template class TopPSamplingLayer; -template class TopPSamplingLayer; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/topPSamplingLayer.h b/cpp/tensorrt_llm/layers/topPSamplingLayer.h deleted file mode 100644 index 974421b0526d..000000000000 --- a/cpp/tensorrt_llm/layers/topPSamplingLayer.h +++ /dev/null @@ -1,72 +0,0 @@ -/* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/layers/baseLayer.h" -#include "tensorrt_llm/runtime/common.h" - -namespace tensorrt_llm::layers -{ - -//! \brief Layer to randomly sample tokens from TopP logits. -//! Layer expects probs precomputed in "logits" tensor -template -class TopPSamplingLayer : public BaseLayer -{ - using Base = BaseLayer; - -public: - TopPSamplingLayer(DecoderDomain const& decoderDomain, std::shared_ptr bufferManager, - bool isDeterministic = true, bool isAirTopP = true); - - void setup(runtime::SizeType32 batchSize, runtime::SizeType32 beamWidth, TensorConstPtr batchSlots, - std::shared_ptr const& setupParams, - std::shared_ptr const& workspace) override; - void forwardAsync(std::shared_ptr const& outputs, - std::shared_ptr const& inputs, - std::shared_ptr const& workspace) override; - - //! @returns workspace needed for this layer in bytes - [[nodiscard]] size_t getWorkspaceSize() const noexcept override; - -protected: - TensorPtr mRuntimeTopKDevice; - TensorPtr mRuntimeTopPDevice; - TensorPtr mInitialTopPDevice; - TensorPtr mTopPDecayDevice; - TensorPtr mTopPMinDevice; - TensorPtr mTopPResetIdsDevice; - - TensorPtr mSkipDecodeDevice; - TensorPtr mSkipDecodeHost; - size_t mWorkspaceSize{0}; - size_t mSetupWorkspaceSize{0}; - - // AirTopP - cudaDeviceProp mDeviceProp; - runtime::SizeType32 mAirTopPBlockNum{0}; - bool mIsDeterministic{true}; - bool mIsAirTopP{false}; - - using Base::mDecoderDomain; - -private: - void allocateBuffer(runtime::SizeType32 batchSize); -}; - -} // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index 5dc1de84a308..c68f75b5f435 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -6,7 +6,6 @@ set(TRTLLM_NB_MODULE set(SRCS batch_manager/algorithms.cpp batch_manager/bindings.cpp - batch_manager/buffers.cpp batch_manager/cacheTransceiver.cpp batch_manager/kvCacheConnector.cpp batch_manager/kvCacheManager.cpp diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp index c13466565342..b49d0ed2dc0a 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp @@ -20,16 +20,13 @@ #include "tensorrt_llm/batch_manager/allocateKvCache.h" #include "tensorrt_llm/batch_manager/assignReqSeqSlots.h" #include "tensorrt_llm/batch_manager/capacityScheduler.h" -#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/medusaBuffers.h" #include "tensorrt_llm/batch_manager/microBatchScheduler.h" #include "tensorrt_llm/batch_manager/pauseRequests.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/common/customCasters.h" -#include "tensorrt_llm/runtime/decoderState.h" #include "tensorrt_llm/runtime/torch.h" #include "tensorrt_llm/runtime/torchView.h" @@ -128,29 +125,4 @@ void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_ nb::arg("generation_requests"), nb::arg("model_config"), nb::arg("cross_kv_cache_manager") = std::nullopt, nb::call_guard()) .def("name", [](AllocateKvCache const&) { return AllocateKvCache::name; }); - - nb::class_(m, CreateNewDecoderRequests::name) - .def(nb::init(), nb::arg("speculative_decoding_fast_logits"), - nb::arg("is_leader_in_orch_mode"), nb::arg("is_normalize_log_probs")) - .def( - "__call__", - [](CreateNewDecoderRequests& self, tr::ModelConfig const& modelConfig, tr::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - tensorrt_llm::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, 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("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/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 6da2739ed681..de43fce2d68a 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -19,14 +19,12 @@ #include "tensorrt_llm/nanobind/common/customCasters.h" #include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/batch_manager/microBatchScheduler.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/batch_manager/sequenceSlotManager.h" #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/common/bindTypes.h" -#include "tensorrt_llm/runtime/gptDecoderBatched.h" #include "tensorrt_llm/runtime/runtimeKernels.h" #include "tensorrt_llm/runtime/torch.h" #include "tensorrt_llm/runtime/torchView.h" @@ -740,82 +738,6 @@ void initBindings(nb::module_& m) nb::arg("encoder_kv_lengths"), nb::arg("previous_batch_indices"), nb::arg("position_id_offset") = 0, nb::call_guard(), "Prepare the persistent CPU input buffers for a simple encoder-decoder batch."); - - m.def( - "make_decoding_batch_input", - [](tb::DecoderInputBuffers& decoderInputBuffers, runtime::decoder::DecoderState& decoderState, - std::vector> const& contextRequests, - std::vector> const& genRequests, tr::ITensor::SharedPtr const& logits, - int beamWidth, std::vector const& numContextLogitsPrefixSum, tr::BufferManager const& manager) - { - std::vector activeSlots; - std::vector generationSteps; - std::vector> logitsVec = {{}}; - - for (int i = 0; i < contextRequests.size(); ++i) - { - if (contextRequests[i]->isLastContextChunk()) - { - activeSlots.push_back(*contextRequests[i]->mSeqSlot); - generationSteps.push_back(contextRequests[i]->getDecodingIter()); - auto contextLogitsOffset = numContextLogitsPrefixSum[i + 1] - 1; - tr::ITensor::SharedPtr logitsView = ITensor::slice(logits, contextLogitsOffset, 1); - - if (beamWidth > 1) - { - // Tile logits of context requests - auto const logitsShape = logitsView->getShape(); - auto const logitsType = logitsView->getDataType(); - auto decoderLogits = manager.gpu(ITensor::makeShape({beamWidth, logitsShape.d[1]}), logitsType); - tensorrt_llm::runtime::kernels::tileTensor( - *decoderLogits, *logitsView, beamWidth, manager.getStream()); - decoderLogits->unsqueeze(0); - logitsVec[0].push_back(std::move(decoderLogits)); - } - else - { - logitsView->unsqueeze(1); - logitsVec[0].push_back(std::move(logitsView)); - } - } - } - - auto genLogitsOffset = numContextLogitsPrefixSum.back(); - for (int i = 0; i < genRequests.size(); ++i) - { - if (genRequests[i]->isGenerationInProgressState()) - { - activeSlots.push_back(*genRequests[i]->mSeqSlot); - generationSteps.push_back(genRequests[i]->getDecodingIter()); - - auto logitsOffset = genLogitsOffset + i * beamWidth; - auto numberOfLogits = beamWidth; - tr::ITensor::SharedPtr logitsView = ITensor::slice(logits, logitsOffset, numberOfLogits); - logitsView->unsqueeze(0); - logitsVec[0].push_back(std::move(logitsView)); - } - } - - auto& batchSlots = decoderInputBuffers.forwardBatchSlots; - batchSlots[0]->resize(activeSlots.size()); - auto batchSlotsRange = tr::BufferRange(*batchSlots[0]); - for (int i = 0; i < activeSlots.size(); ++i) - { - batchSlotsRange[i] = activeSlots[i]; - } - - decoderInputBuffers.batchLogits = logitsVec; - - auto const maxBeamWidth = decoderState.getMaxBeamWidth(); - if (maxBeamWidth > 1) - { - // For Variable-Beam-Width-Search - decoderState.getJointDecodingInput().generationSteps = generationSteps; - } - }, - nb::arg("decoder_input_buffers"), nb::arg("decoder_state"), nb::arg("context_requests"), - nb::arg("generation_requests"), nb::arg("logits"), nb::arg("beam_width"), - nb::arg("num_context_logits_prefix_sum"), nb::arg("buffer_manager"), "Make decoding batch input."); } } // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/buffers.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/buffers.cpp deleted file mode 100644 index 9b8e441745cc..000000000000 --- a/cpp/tensorrt_llm/nanobind/batch_manager/buffers.cpp +++ /dev/null @@ -1,74 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "buffers.h" - -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/nanobind/batch_manager/llmRequest.h" -#include "tensorrt_llm/nanobind/common/customCasters.h" -#include "tensorrt_llm/runtime/torch.h" - -#include -#include -#include -#include - -namespace nb = nanobind; -namespace tb = tensorrt_llm::batch_manager; -namespace tr = tensorrt_llm::runtime; - -using tr::SizeType32; - -namespace tensorrt_llm::nanobind::batch_manager -{ - -void Buffers::initBindings(nb::module_& m) -{ - nb::class_(m, "DecoderInputBuffers") - .def(nb::init(), nb::arg("max_batch_size"), - nb::arg("max_tokens_per_engine_step"), nb::arg("manager")) - .def_rw("setup_batch_slots", &tb::DecoderInputBuffers::setupBatchSlots) - .def_rw("setup_batch_slots_device", &tb::DecoderInputBuffers::setupBatchSlotsDevice) - .def_rw("fill_values", &tb::DecoderInputBuffers::fillValues) - .def_rw("fill_values_device", &tb::DecoderInputBuffers::fillValuesDevice) - .def_rw("inputs_ids", &tb::DecoderInputBuffers::inputsIds) - .def_rw("forward_batch_slots", &tb::DecoderInputBuffers::forwardBatchSlots) - .def_rw("decoder_logits", &tb::DecoderInputBuffers::decoderLogits) - .def_rw("decoder_requests", &tb::DecoderInputBuffers::decoderRequests); - - nb::class_(m, "DecoderOutputBuffers") - .def_rw("sequence_lengths_host", &tb::DecoderOutputBuffers::sequenceLengthsHost) - .def_rw("finished_sum_host", &tb::DecoderOutputBuffers::finishedSumHost) - .def_prop_ro("new_output_tokens_host", - [](tb::DecoderOutputBuffers& self) { return tr::Torch::tensor(self.newOutputTokensHost); }) - .def_rw("cum_log_probs_host", &tb::DecoderOutputBuffers::cumLogProbsHost) - .def_rw("log_probs_host", &tb::DecoderOutputBuffers::logProbsHost) - .def_rw("finish_reasons_host", &tb::DecoderOutputBuffers::finishReasonsHost); - - nb::class_(m, "SlotDecoderBuffers") - .def(nb::init(), - nb::arg("max_beam_width"), nb::arg("max_seq_len"), nb::arg("buffer_manager")) - .def_rw("output_ids", &tb::SlotDecoderBuffers::outputIds) - .def_rw("output_ids_host", &tb::SlotDecoderBuffers::outputIdsHost) - .def_rw("sequence_lengths_host", &tb::SlotDecoderBuffers::sequenceLengthsHost) - .def_rw("cum_log_probs", &tb::SlotDecoderBuffers::cumLogProbs) - .def_rw("cum_log_probs_host", &tb::SlotDecoderBuffers::cumLogProbsHost) - .def_rw("log_probs", &tb::SlotDecoderBuffers::logProbs) - .def_rw("log_probs_host", &tb::SlotDecoderBuffers::logProbsHost) - .def_rw("finish_reasons_host", &tb::SlotDecoderBuffers::finishReasonsHost); -} -} // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/buffers.h b/cpp/tensorrt_llm/nanobind/batch_manager/buffers.h deleted file mode 100644 index d33570f7d36e..000000000000 --- a/cpp/tensorrt_llm/nanobind/batch_manager/buffers.h +++ /dev/null @@ -1,30 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include -namespace nb = nanobind; - -namespace tensorrt_llm::nanobind::batch_manager -{ -class Buffers -{ -public: - static void initBindings(nb::module_& m); -}; -} // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index a2054dbd7217..715d057f7edc 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -35,7 +35,6 @@ #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/batch_manager/algorithms.h" #include "tensorrt_llm/nanobind/batch_manager/bindings.h" -#include "tensorrt_llm/nanobind/batch_manager/buffers.h" #include "tensorrt_llm/nanobind/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/nanobind/batch_manager/kvCacheConnector.h" #include "tensorrt_llm/nanobind/batch_manager/kvCacheManager.h" @@ -73,14 +72,6 @@ using OptVec = std::optional>; #error "TRTLLM_NB_MODULE must be defined" #endif -namespace -{ -tr::SamplingConfig makeSamplingConfig(std::vector const& configs) -{ - return tr::SamplingConfig(configs); -} -} // namespace - NB_MODULE(TRTLLM_NB_MODULE, m) { m.doc() = "TensorRT LLM Python bindings for C++ runtime"; @@ -458,10 +449,6 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def("__setstate__", SamplingConfigSetState) .def("__eq__", &tr::SamplingConfig::operator==); - nb::bind_vector>(m, "SamplingConfigVector"); - - m.def("make_sampling_config", &makeSamplingConfig, nb::arg("configs")); - nb::class_(m, "GptJsonConfig") .def(nb::init>(), @@ -514,7 +501,6 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def_prop_ro("uvm", &tr::MemoryCounters::getUVM); tensorrt_llm::nanobind::process_group::initBindings(mInternalProcessGroup); - tpb::Buffers::initBindings(mInternalBatchManager); tensorrt_llm::nanobind::runtime::initBindings(mInternalRuntime); tensorrt_llm::nanobind::testing::initKvCacheTestUtilBindings(mInternalTesting); tpb::initBindings(mInternalBatchManager); diff --git a/cpp/tensorrt_llm/nanobind/common/customCasters.h b/cpp/tensorrt_llm/nanobind/common/customCasters.h index 8c202b9387fa..d118abfac7fd 100644 --- a/cpp/tensorrt_llm/nanobind/common/customCasters.h +++ b/cpp/tensorrt_llm/nanobind/common/customCasters.h @@ -18,7 +18,6 @@ #pragma once #include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/common/optionalRef.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/samplingConfig.h" @@ -44,8 +43,6 @@ // Opaque bindings NB_MAKE_OPAQUE(tensorrt_llm::batch_manager::ReqIdsSet) -NB_MAKE_OPAQUE(std::vector) -NB_MAKE_OPAQUE(std::vector) namespace nb = nanobind; diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp index eec3cd79bac1..728857ebb55c 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -27,13 +27,7 @@ #include "tensorrt_llm/nanobind/common/customCasters.h" #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/decodingInput.h" -#include "tensorrt_llm/runtime/decodingOutput.h" -#include "tensorrt_llm/runtime/gptDecoder.h" -#include "tensorrt_llm/runtime/gptDecoderBatched.h" #include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iGptDecoderBatched.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/ipcUtils.h" #include "tensorrt_llm/runtime/lookaheadBuffers.h" @@ -60,44 +54,6 @@ namespace tr = tensorrt_llm::runtime; namespace te = tensorrt_llm::executor; -class PyIGptDecoder : public tr::IGptDecoder -{ -public: - NB_TRAMPOLINE(tr::IGptDecoder, 5); - - void setup(tr::SamplingConfig const& samplingConfig, size_t batchSize, - tr::DecodingInput::TensorConstPtr const& batchSlots, - std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, - std::optional> const& lookaheadPrompt = std::nullopt, - std::optional> const& lookaheadAlgoConfigs = std::nullopt) override - { - NB_OVERRIDE_PURE(setup, samplingConfig, batchSize, batchSlots, output, explicitDraftTokensDType, - lookaheadPrompt, lookaheadAlgoConfigs); - } - - void forwardAsync(tr::DecodingOutput& output, tr::DecodingInput const& input) override - { - NB_OVERRIDE_PURE(forwardAsync, output, input); - } - - void forwardSync(tr::DecodingOutput& output, tr::DecodingInput const& input) override - { - NB_OVERRIDE_PURE(forwardSync, output, input); - } - - tr::SamplingConfig const& getSamplingConfig() override - { - NB_OVERRIDE_PURE(getSamplingConfig); - } - - void disableLookahead(std::optional const& samplingConfig, tr::SizeType32 batchSize, - tr::DecodingInput::TensorConstPtr batchSlots) override - { - NB_OVERRIDE_PURE(disableLookahead, samplingConfig, batchSize, batchSlots); - } -}; - namespace tensorrt_llm::nanobind::runtime { @@ -133,124 +89,11 @@ void initBindings(nb::module_& m) .def_rw("packed_masks", &tr::LookaheadDecodingBuffers::packedMasks) .def_rw("position_ids", &tr::LookaheadDecodingBuffers::positionIds); - nb::class_(m, "ExplicitDraftTokensBuffersInputs") - .def("create", &tr::ExplicitDraftTokensBuffers::Inputs::create, nb::arg("max_num_sequences"), - nb::arg("runtime"), nb::arg("model_config"), nb::arg("world_config"), - nb::call_guard()) - .def_rw("temperatures", &tr::ExplicitDraftTokensBuffers::Inputs::temperatures) - .def_rw("position_ids_base", &tr::ExplicitDraftTokensBuffers::Inputs::positionIdsBase) - .def_rw("generation_lengths", &tr::ExplicitDraftTokensBuffers::Inputs::generationLengths) - .def_rw("random_data_sample", &tr::ExplicitDraftTokensBuffers::Inputs::randomDataSample) - .def_rw("random_data_validation", &tr::ExplicitDraftTokensBuffers::Inputs::randomDataValidation) - .def_rw("draft_tokens", &tr::ExplicitDraftTokensBuffers::Inputs::draftTokens) - .def_rw("draft_indices", &tr::ExplicitDraftTokensBuffers::Inputs::draftIndices) - .def_rw("draft_probs", &tr::ExplicitDraftTokensBuffers::Inputs::draftProbs) - .def_rw("packed_masks", &tr::ExplicitDraftTokensBuffers::Inputs::packedMasks) - .def_rw("position_ids", &tr::ExplicitDraftTokensBuffers::Inputs::positionIds) - .def_rw("max_gen_length_host", &tr::ExplicitDraftTokensBuffers::Inputs::maxGenLengthHost) - .def_rw("generation_lengths_host", &tr::ExplicitDraftTokensBuffers::Inputs::generationLengthsHost); - - nb::class_(m, "DecodingInput"); - nb::class_(m, "DecodingOutput"); - nb::class_(m, "CudaEvent") .def(nb::init(), nb::arg("flags") = cudaEventDisableTiming, nb::call_guard()) .def("synchronize", &tr::CudaEvent::synchronize, nb::call_guard()); - nb::class_(m, "IGptDecoder") - .def( - "setup", - [](tr::IGptDecoder& self, tr::SamplingConfig const& samplingConfig, size_t batchSize, - at::Tensor const& batchSlots, std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, - std::optional> const& lookaheadPrompt = std::nullopt, - std::optional> const& lookaheadAlgoConfigs = std::nullopt) - { - auto tensorPtrBatchSlots = tr::TorchView::of(batchSlots); - self.setup(samplingConfig, batchSize, std::move(tensorPtrBatchSlots), output, explicitDraftTokensDType, - lookaheadPrompt, lookaheadAlgoConfigs); - }, - nb::arg("sampling_config"), nb::arg("batch_size"), nb::arg("batch_slots"), nb::arg("output") = std::nullopt, - nb::arg("explicit_draft_tokens_d_type") = std::nullopt, nb::arg("lookahead_prompt") = std::nullopt, - nb::arg("lookahead_algo_configs") = std::nullopt, nb::call_guard()); - - nb::class_(m, "DecoderState") - .def(nb::init<>(), nb::call_guard()) - .def("setup", &tr::decoder::DecoderState::setup, nb::arg("max_num_sequences"), nb::arg("max_beam_width"), - nb::arg("max_attention_window"), nb::arg("sink_token_length"), nb::arg("max_sequence_length"), - nb::arg("dtype"), nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager"), - nb::call_guard()) - .def("setup_cache_indirection", &tr::decoder::DecoderState::setupCacheIndirection, nb::arg("max_num_sequences"), - nb::arg("max_beam_width"), nb::arg("max_attention_window"), nb::arg("buffer_manager"), - nb::call_guard()) - .def("setup_speculative_decoding", &tr::decoder::DecoderState::setupSpeculativeDecoding, - nb::arg("speculative_decoding_mode"), nb::arg("max_tokens_per_engine_step"), nb::arg("dtype"), - nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager"), - nb::call_guard()) - .def_prop_ro("joint_decoding_input", &tr::decoder::DecoderState::getJointDecodingInput) - .def_prop_ro("joint_decoding_output", &tr::decoder::DecoderState::getJointDecodingOutput) - .def_prop_ro("cache_indirection_input", &tr::decoder::DecoderState::getCacheIndirectionInput) - .def_prop_ro("cache_indirection_output", &tr::decoder::DecoderState::getCacheIndirectionOutput) - .def_prop_ro( - "sequence_lengths", nb::overload_cast<>(&tr::decoder::DecoderState::getSequenceLengths, nb::const_)) - .def("get_sequence_lengths", - nb::overload_cast(&tr::decoder::DecoderState::getSequenceLengths, nb::const_), - nb::arg("batch_idx"), nb::call_guard()) - .def_prop_ro("all_new_tokens", &tr::decoder::DecoderState::getAllNewTokens) - .def_prop_ro("finished_sum", &tr::decoder::DecoderState::getFinishedSum) - .def_prop_ro("finish_reasons", &tr::decoder::DecoderState::getFinishReasons) - .def_prop_ro("ids", nb::overload_cast<>(&tr::decoder::DecoderState::getIds, nb::const_)) - .def("get_ids", nb::overload_cast(&tr::decoder::DecoderState::getIds, nb::const_), - nb::arg("batch_idx"), nb::call_guard()) - .def_prop_ro("gathered_ids", nb::overload_cast<>(&tr::decoder::DecoderState::getGatheredIds, nb::const_)) - .def("get_gathered_ids", - nb::overload_cast(&tr::decoder::DecoderState::getGatheredIds, nb::const_), - nb::arg("batch_idx"), nb::call_guard()) - .def_prop_ro("parent_ids", &tr::decoder::DecoderState::getParentIds) - .def_prop_ro("cum_log_probs", nb::overload_cast<>(&tr::decoder::DecoderState::getCumLogProbs, nb::const_)) - .def("get_cum_log_probs", - nb::overload_cast(&tr::decoder::DecoderState::getCumLogProbs, nb::const_), - nb::arg("batch_idx"), nb::call_guard()) - .def_prop_ro("log_probs", nb::overload_cast<>(&tr::decoder::DecoderState::getLogProbs, nb::const_)) - .def("get_log_probs", nb::overload_cast(&tr::decoder::DecoderState::getLogProbs, nb::const_), - nb::arg("batch_idx"), nb::call_guard()) - .def_prop_ro("next_draft_tokens", &tr::decoder::DecoderState::getNextDraftTokens) - .def_prop_ro("prev_draft_tokens_lengths", &tr::decoder::DecoderState::getPrevDraftTokensLengths) - .def_prop_ro("next_draft_tokens_lengths", &tr::decoder::DecoderState::getNextDraftTokensLengths) - .def_prop_ro("accepted_lengths_cum_sum", &tr::decoder::DecoderState::getAcceptedLengthsCumSum) - .def_prop_ro("accepted_packed_paths", &tr::decoder::DecoderState::getAcceptedPackedPaths) - .def_prop_ro("max_beam_width", &tr::decoder::DecoderState::getMaxBeamWidth) - .def_prop_ro("max_sequence_length", &tr::decoder::DecoderState::getMaxSequenceLength) - .def_prop_ro("max_decoding_decoder_tokens", &tr::decoder::DecoderState::getMaxDecodingDecoderTokens) - .def_prop_ro("max_decoding_engine_tokens", &tr::decoder::DecoderState::getMaxDecodingEngineTokens) - .def_prop_ro("num_decoding_engine_tokens", - nb::overload_cast<>(&tr::decoder::DecoderState::getNumDecodingEngineTokens, nb::const_)) - .def("get_num_decoding_engine_tokens", - nb::overload_cast(&tr::decoder::DecoderState::getNumDecodingEngineTokens, nb::const_), - nb::arg("batch_idx"), nb::call_guard()) - .def("set_num_decoding_engine_tokens", &tr::decoder::DecoderState::setNumDecodingEngineTokens, - nb::arg("batch_idx"), nb::arg("num_tokens"), nb::call_guard()) - .def_prop_ro("speculative_decoding_mode", &tr::decoder::DecoderState::getSpeculativeDecodingMode) - .def_prop_rw("generation_steps", &tr::decoder::DecoderState::getGenerationSteps, - &tr::decoder::DecoderState::setGenerationSteps); - - nb::class_(m, "GptDecoderBatched") - .def(nb::init(), nb::arg("stream"), - nb::call_guard()) - .def("setup", &tr::GptDecoderBatched::setup, nb::arg("mode"), nb::arg("max_num_sequences"), - nb::arg("max_beam_width"), nb::arg("dtype"), nb::arg("model_config"), nb::arg("world_config"), - nb::call_guard()) - .def("forward_async", &tr::GptDecoderBatched::forwardAsync, nb::arg("decoder_state"), nb::arg("input"), - nb::call_guard()) - .def("underlying_decoder", &tr::GptDecoderBatched::getUnderlyingDecoder, nb::rv_policy::reference) - .def("finalize", &tr::GptDecoderBatched::finalize, nb::arg("decoder_state"), nb::arg("batch_idx"), - nb::arg("sampling_config"), nb::arg("streaming"), nb::call_guard()) - .def_prop_ro( - "decoder_stream", - [](tr::GptDecoderBatched& self) -> tr::CudaStream const& { return *self.getDecoderStream(); }, - nb::rv_policy::reference); - m.def( "lamport_initialize_all", [](intptr_t buffer_0, intptr_t buffer_1, intptr_t buffer_2, size_t size) diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index 11a9391c0e69..0bd18753fc76 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -22,18 +22,11 @@ set(SRCS utils/speculativeChoicesUtils.cpp bufferManager.cpp cudaMemPool.cpp - decodingLayerWorkspace.cpp - eagleBuffers.cpp - explicitDraftTokensBuffers.cpp lookaheadBuffers.cpp loraManager.cpp loraUtils.cpp loraModule.cpp loraCache.cpp - decodingOutput.cpp - decoderState.cpp - gptDecoder.cpp - gptDecoderBatched.cpp gptJsonConfig.cpp iBuffer.cpp iTensor.cpp diff --git a/cpp/tensorrt_llm/runtime/decoderState.cpp b/cpp/tensorrt_llm/runtime/decoderState.cpp deleted file mode 100644 index 83037b2431cb..000000000000 --- a/cpp/tensorrt_llm/runtime/decoderState.cpp +++ /dev/null @@ -1,671 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -namespace tk = tensorrt_llm::kernels; - -namespace tensorrt_llm::runtime::decoder -{ -using TensorPtr = DecoderState::TensorPtr; - -BeamSearchBuffers::BeamSearchBuffers(BufferManager const& bufferManager) - : mOutputBeamHypotheses{} - , mCumLogProbsTmp(bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT)) -{ - mOutputBeamHypotheses.empty(bufferManager); - mCumLogProbsTmp = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); - - int device; - cudaGetDevice(&device); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, device); - mNumSMs = deviceProp.multiProcessorCount; -} - -void BeamSearchBuffers::reshape(SizeType32 maxBeamWidth, SizeType32 maxSequenceLength) -{ - mOutputBeamHypotheses.reshape(1, maxBeamWidth, maxSequenceLength); - mCumLogProbsTmp->reshape(ITensor::makeShape({1, maxBeamWidth})); -} - -DecoderState::DecoderState() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mJointDecodingInput = std::make_unique(); - mJointDecodingOutput = std::make_unique(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, tensorrt_llm::DataType dtype, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - setupBuffers(dtype, bufferManager); - reshapeBuffers(maxNumSequences, maxBeamWidth, maxAttentionWindow, sinkTokenLength, maxSequenceLength, modelConfig, - worldConfig, bufferManager); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::setupBuffers(tensorrt_llm::DataType dtype, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - auto constexpr nvFloatType = TRTDataType::value; - - auto& dInput = mJointDecodingInput; - TLLM_CHECK(static_cast(dInput)); - dInput->endIds = bufferManager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - dInput->batchSlots = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvSizeType); - - dInput->sequenceLimitLength = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - dInput->lengths = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - - auto& dOutput = mJointDecodingOutput; - TLLM_CHECK(static_cast(dOutput)); - dOutput->ids = bufferManager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - dOutput->gatheredIds = bufferManager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - - dOutput->newTokensSteps = bufferManager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - dOutput->parentIds = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - - dOutput->lengths = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - - dOutput->finishedSum = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - dOutput->cumLogProbs = bufferManager.emptyTensor(MemoryType::kGPU, nvFloatType); - dOutput->logProbs = bufferManager.emptyTensor(MemoryType::kGPU, nvFloatType); - dOutput->beamHypotheses.empty(bufferManager); - - dOutput->finishReasons - = bufferManager.emptyTensor(MemoryType::kGPU, TRTDataType::value); - dInput->finishReasons = dOutput->finishReasons; - - dOutput->logProbsTiled = bufferManager.emptyTensor(MemoryType::kGPU, nvFloatType); - - dInput->stopWordsPtrs = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, TRTDataType::value); - dInput->stopWordsLens = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvSizeType); - dInput->badWordsPtrs = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, TRTDataType::value); - dInput->badWordsLens = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvSizeType); - dInput->embeddingBias = bufferManager.emptyTensor(MemoryType::kGPU, dtype); - - mBeamSearchBuffers = std::make_unique(bufferManager); - - setupCacheIndirectionBuffers(bufferManager); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::setupSpeculativeDecoding(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - setupSpeculativeDecodingBuffers(speculativeDecodingMode, dtype, bufferManager); - reshapeSpeculativeDecodingBuffers( - speculativeDecodingMode, maxTokensPerEngineStep, modelConfig, worldConfig, bufferManager); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::setupSpeculativeDecodingBuffers(SpeculativeDecodingMode const speculativeDecodingMode, - tensorrt_llm::DataType dtype, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mSpeculativeDecodingMode = speculativeDecodingMode; - - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - - auto& dInput = mJointDecodingInput; - auto& dOutput = mJointDecodingOutput; - - if (speculativeDecodingMode.isMedusa()) - { - DecodingInput::MedusaInputs medusaInputs; - medusaInputs.medusaPaths = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - medusaInputs.medusaTreeIds = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - medusaInputs.medusaCurTokensPerStep = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - medusaInputs.medusaTargetTokensPerStep = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - dInput->medusaInputs = medusaInputs; - } - - DecodingOutput::SpeculativeDecodingOutputs speculativeDecodingOutputs; - if (speculativeDecodingMode.predictsDraftTokens()) - { - speculativeDecodingOutputs.nextDraftTokens - = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - if (speculativeDecodingMode.variableDraftLength()) - { - speculativeDecodingOutputs.nextDraftTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - speculativeDecodingOutputs.prevDraftTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - } - } - if (speculativeDecodingMode.isLookaheadDecoding()) - { - dInput->lookaheadInputs = DecodingInput::LookaheadInputs(); - } - if (speculativeDecodingMode.needsKVCacheRewind()) - { - speculativeDecodingOutputs.acceptedTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - speculativeDecodingOutputs.acceptedLengthsCumSum - = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - speculativeDecodingOutputs.pathsOffsets - = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - } - dOutput->speculativeDecodingOutputs = speculativeDecodingOutputs; - - if (speculativeDecodingMode.isDraftTokensExternal()) - { - DecodingInput::ExternalDraftTokensInputs externalDraftTokensInputs; - - externalDraftTokensInputs.draftLogits = bufferManager.emptyTensor(MemoryType::kGPU, dtype); - externalDraftTokensInputs.draftLogitsHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, dtype); - externalDraftTokensInputs.draftProbs = bufferManager.emptyTensor(MemoryType::kGPU, dtype); - externalDraftTokensInputs.targetProbs = bufferManager.emptyTensor(MemoryType::kGPU, dtype); - externalDraftTokensInputs.numDraftTokens = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - externalDraftTokensInputs.numDraftTokensHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvSizeType); - externalDraftTokensInputs.useDraftLogits - = bufferManager.emptyTensor(MemoryType::kGPU, TRTDataType::value); - externalDraftTokensInputs.useDraftLogitsHost - = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, TRTDataType::value); - externalDraftTokensInputs.draftTokenIds = bufferManager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - externalDraftTokensInputs.draftTokenIdsHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvTokenIdType); - - dInput->externalDraftTokensInputs = externalDraftTokensInputs; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::reshapeBuffers(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& stream = bufferManager.getStream(); - - TLLM_CHECK(maxNumSequences > 0); - TLLM_CHECK(maxBeamWidth > 0); - TLLM_CHECK(mMaxDecodingEngineTokens > 0); - TLLM_CHECK(maxSequenceLength > 0); - mMaxNumSequences = maxNumSequences; - mMaxBeamWidth = maxBeamWidth; - mMaxSequenceLength = maxSequenceLength; - - mNumDecodingEngineTokens.clear(); - mNumDecodingEngineTokens.resize(mMaxNumSequences, 0); - - // setup input - auto& dInput = *mJointDecodingInput; - dInput.maxLength = mMaxSequenceLength; - dInput.maxAttentionWindow = maxAttentionWindow; - dInput.sinkTokenLength = sinkTokenLength; - dInput.stopWordsLists.resize(mMaxNumSequences); - dInput.badWordsLists.resize(mMaxNumSequences); - - auto const maxNumSequencesShape = ITensor::makeShape({mMaxNumSequences}); - auto const maxNumSequencesXmaxBeamWidthShape = ITensor::makeShape({mMaxNumSequences, mMaxBeamWidth}); - - const_cast(*dInput.endIds).reshape(maxNumSequencesShape); - auto& sequenceLimitLength = const_cast(*dInput.sequenceLimitLength); - sequenceLimitLength.reshape(maxNumSequencesShape); - kernels::invokeFill(sequenceLimitLength, mMaxSequenceLength, stream); - auto& inputLengths = const_cast(*dInput.lengths); - inputLengths.reshape(maxNumSequencesXmaxBeamWidthShape); - bufferManager.setZero(inputLengths); - - dInput.beamWidths.clear(); - dInput.beamWidths.resize(mMaxNumSequences, 0); - - auto const maxTotalTokensShape = ITensor::makeShape({mMaxNumSequences, mMaxBeamWidth, mMaxSequenceLength}); - - // setup output - auto& dOutput = *mJointDecodingOutput; - dOutput.ids->reshape(maxTotalTokensShape); - - auto const maxNewTokensShape = ITensor::makeShape({mMaxDecodingEngineTokens, mMaxNumSequences, mMaxBeamWidth}); - - dOutput.finishReasons->reshape(maxNumSequencesXmaxBeamWidthShape); - bufferManager.setZero(*dOutput.finishReasons); - - dOutput.parentIds->reshape(maxTotalTokensShape); - - dOutput.lengths->reshape(maxNumSequencesXmaxBeamWidthShape); - bufferManager.setZero(*dOutput.lengths); - - dOutput.finishedSum->reshape(maxNumSequencesShape); - bufferManager.setZero(*dOutput.finishedSum); - - dOutput.newTokensSteps->reshape(maxNewTokensShape); - bufferManager.setZero(*dOutput.newTokensSteps); - - dOutput.cumLogProbs->reshape(maxNumSequencesXmaxBeamWidthShape); - bufferManager.setZero(*dOutput.cumLogProbs); - - dOutput.logProbs->reshape(maxTotalTokensShape); - bufferManager.setZero(*dOutput.logProbs); - - dOutput.logProbsTiled->reshape(ITensor::makeShape({mMaxSequenceLength, mMaxNumSequences, mMaxBeamWidth})); - bufferManager.setZero(*dOutput.logProbsTiled); - - if (mMaxBeamWidth > 1) - { - dOutput.beamHypotheses.reshape(mMaxNumSequences, mMaxBeamWidth, mMaxSequenceLength); - mBeamSearchBuffers->reshape(mMaxBeamWidth, mMaxSequenceLength); - - reshapeCacheIndirectionBuffers(mMaxNumSequences, mMaxBeamWidth, maxAttentionWindow); - - dOutput.gatheredIds->reshape(maxTotalTokensShape); - } - else - { - dOutput.gatheredIds = dOutput.ids; - } - - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - const_cast(*dInput.embeddingBias) - .reshape(ITensor::makeShape({mMaxNumSequences, static_cast(vocabSizePadded)})); - const_cast(*dInput.badWordsPtrs).reshape(maxNumSequencesShape); - const_cast(*dInput.badWordsLens).reshape(maxNumSequencesShape); - const_cast(*dInput.stopWordsPtrs).reshape(maxNumSequencesShape); - const_cast(*dInput.stopWordsLens).reshape(maxNumSequencesShape); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::setupCacheIndirection(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - SizeType32 maxAttentionWindow, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - setupCacheIndirectionBuffers(bufferManager); - reshapeCacheIndirectionBuffers(maxNumSequences, maxBeamWidth, maxAttentionWindow); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::setupCacheIndirectionBuffers(BufferManager const& bufferManager) -{ - auto constexpr nvSizeType = TRTDataType::value; - mJointDecodingInput->cacheIndirection = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); - mJointDecodingOutput->cacheIndirection = bufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); -} - -void DecoderState::reshapeCacheIndirectionBuffers( - SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow) -{ - mJointDecodingInput->cacheIndirection->reshape( - ITensor::makeShape({maxNumSequences, maxBeamWidth, maxAttentionWindow})); - mJointDecodingOutput->cacheIndirection->reshape( - ITensor::makeShape({maxNumSequences, maxBeamWidth, maxAttentionWindow})); -} - -void DecoderState::reshapeSpeculativeDecodingBuffers(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto& dInput = *mJointDecodingInput; - auto& dOutput = *mJointDecodingOutput; - - TLLM_CHECK(maxTokensPerEngineStep > 0); - mMaxDecodingEngineTokens = maxTokensPerEngineStep; - - TLLM_CHECK_WITH_INFO((mMaxDecodingEngineTokens == 1 && speculativeDecodingMode.isNone()) - || (mMaxDecodingEngineTokens > 1 && !speculativeDecodingMode.isNone()), - "Max tokens per engine step is %d, but must be equal to 1 when no speculative decoding is configured, " - "or > 1 for any speculative decoding mode.", - mMaxDecodingEngineTokens); - - auto const maxNewTokensShape = ITensor::makeShape({mMaxDecodingEngineTokens, mMaxNumSequences, mMaxBeamWidth}); - dOutput.newTokensSteps->reshape(maxNewTokensShape); - bufferManager.setZero(*dOutput.newTokensSteps); - - if (speculativeDecodingMode.predictsDraftTokens()) - { - mMaxDecodingDecoderTokens = mMaxDecodingEngineTokens; - } - else - { - mMaxDecodingDecoderTokens = 1; - } - - if (speculativeDecodingMode.isNone()) - { - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return; - } - - auto const maxNumSequencesShape = ITensor::makeShape({mMaxNumSequences}); - - if (speculativeDecodingMode.isDraftTokensExternal()) - { - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const probsShape = ITensor::makeShape( - {mMaxNumSequences, mMaxDecodingEngineTokens, mMaxBeamWidth, static_cast(vocabSizePadded)}); - dInput.externalDraftTokensInputs->draftProbs->reshape(probsShape); - dInput.externalDraftTokensInputs->targetProbs->reshape(probsShape); - - auto const logitsShape = ITensor::makeShape( - {mMaxNumSequences, mMaxDecodingEngineTokens, static_cast(vocabSizePadded)}); - dInput.externalDraftTokensInputs->draftLogits->reshape(logitsShape); - dInput.externalDraftTokensInputs->draftLogitsHost->reshape(logitsShape); - - auto const tokenIdsShape = ITensor::makeShape({mMaxNumSequences, mMaxDecodingEngineTokens}); - dInput.externalDraftTokensInputs->draftTokenIds->reshape(tokenIdsShape); - dInput.externalDraftTokensInputs->draftTokenIdsHost->reshape(tokenIdsShape); - - dInput.externalDraftTokensInputs->numDraftTokens->reshape(maxNumSequencesShape); - dInput.externalDraftTokensInputs->numDraftTokensHost->reshape(maxNumSequencesShape); - dInput.externalDraftTokensInputs->useDraftLogits->reshape(maxNumSequencesShape); - dInput.externalDraftTokensInputs->useDraftLogitsHost->reshape(maxNumSequencesShape); - } - - if (speculativeDecodingMode.isMedusa()) - { - auto const speculativeDecodingModule = modelConfig.getSpeculativeDecodingModulePtr(); - auto& medusaPaths = const_cast(*dInput.medusaInputs->medusaPaths); - medusaPaths.reshape(ITensor::makeShape({mMaxNumSequences, speculativeDecodingModule->getMaxDecodingTokens(), - speculativeDecodingModule->getMaxPathLen()})); - bufferManager.setMem(medusaPaths, -1); - - auto& medusaTreeIds = const_cast(*dInput.medusaInputs->medusaTreeIds); - medusaTreeIds.reshape( - ITensor::makeShape({mMaxNumSequences, speculativeDecodingModule->getMaxDecodingDraftTokens()})); - bufferManager.setZero(medusaTreeIds); - auto& curTokensPerStep = const_cast(*dInput.medusaInputs->medusaCurTokensPerStep); - auto& targetTokensPerStep = const_cast(*dInput.medusaInputs->medusaTargetTokensPerStep); - curTokensPerStep.reshape(maxNumSequencesShape); - targetTokensPerStep.reshape(maxNumSequencesShape); - bufferManager.setZero(curTokensPerStep); - bufferManager.setZero(targetTokensPerStep); - } - - if (speculativeDecodingMode.predictsDraftTokens()) - { - dOutput.speculativeDecodingOutputs->nextDraftTokens->reshape( - ITensor::makeShape({mMaxNumSequences, mMaxDecodingEngineTokens - 1})); - if (speculativeDecodingMode.variableDraftLength()) - { - dOutput.speculativeDecodingOutputs->nextDraftTokensLen->reshape(maxNumSequencesShape); - dOutput.speculativeDecodingOutputs->prevDraftTokensLen->reshape(maxNumSequencesShape); - } - } - if (speculativeDecodingMode.needsKVCacheRewind()) - { - auto const speculativeDecodingModule = modelConfig.getSpeculativeDecodingModulePtr(); - dOutput.speculativeDecodingOutputs->acceptedTokensLen->reshape(maxNumSequencesShape); - dOutput.speculativeDecodingOutputs->acceptedLengthsCumSum->reshape(ITensor::makeShape({mMaxNumSequences + 1})); - dOutput.speculativeDecodingOutputs->pathsOffsets->reshape( - ITensor::makeShape({mMaxNumSequences * speculativeDecodingModule->getMaxDraftPathLen()})); - } - - if (speculativeDecodingMode.isExplicitDraftTokens()) - { - mJointDecodingOutput->explicitDraftTokensBuffers = runtime::ExplicitDraftTokensBuffers::Inputs(); - mJointDecodingOutput->explicitDraftTokensBuffers->create( - mMaxNumSequences, bufferManager, modelConfig, worldConfig); - } - else if (speculativeDecodingMode.isEagle()) - { - mJointDecodingOutput->eagleBuffers = runtime::EagleBuffers::Inputs(); - mJointDecodingOutput->eagleBuffers->create(mMaxNumSequences, bufferManager, modelConfig, worldConfig); - } - else if (speculativeDecodingMode.isLookaheadDecoding()) - { - mJointDecodingOutput->lookaheadOutputs - = runtime::LookaheadDecodingBuffers(mMaxNumSequences, mMaxDecodingEngineTokens, bufferManager); - mJointDecodingInput->lookaheadInputs->tokensPerStep = mJointDecodingOutput->lookaheadOutputs->generationLengths; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void DecoderState::disableLookahead(RequestVector const& genRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mSpeculativeDecodingMode = SpeculativeDecodingMode::None(); - - mMaxDecodingEngineTokens = 1; - mMaxDecodingDecoderTokens = 1; - mJointDecodingInput->lookaheadInputs.reset(); - - auto const maxNewTokensShape = ITensor::makeShape({mMaxDecodingEngineTokens, mMaxNumSequences, mMaxBeamWidth}); - mJointDecodingOutput->newTokensSteps->reshape(maxNewTokensShape); - - for (auto const& llmReq : genRequests) - { - if (llmReq->mSeqSlot) - { - setNumDecodingEngineTokens(llmReq->mSeqSlot.value(), 1); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -TensorPtr DecoderState::getFinishedSum() const -{ - return mJointDecodingOutput->finishedSum; -} - -TensorPtr DecoderState::getFinishReasons() const -{ - return mJointDecodingOutput->finishReasons; -} - -TensorPtr DecoderState::getIds() const -{ - return mJointDecodingOutput->ids; -} - -TensorPtr DecoderState::getIds(SizeType32 batchIdx) const -{ - return ITensor::at(mJointDecodingOutput->ids, {batchIdx}); -} - -TensorPtr DecoderState::getGatheredIds() const -{ - return mJointDecodingOutput->gatheredIds; -} - -TensorPtr DecoderState::getGatheredIds(SizeType32 batchIdx) const -{ - return ITensor::at(mJointDecodingOutput->gatheredIds, {batchIdx}); -} - -TensorPtr DecoderState::getParentIds() const -{ - return mJointDecodingOutput->parentIds; -} - -TensorPtr DecoderState::getCumLogProbs() const -{ - return mJointDecodingOutput->cumLogProbs; -} - -TensorPtr DecoderState::getCumLogProbs(SizeType32 batchIdx) const -{ - return ITensor::at(mJointDecodingOutput->cumLogProbs, {batchIdx}); -} - -TensorPtr DecoderState::getLogProbs() const -{ - return mJointDecodingOutput->logProbs; -} - -TensorPtr DecoderState::getLogProbs(SizeType32 batchIdx) const -{ - return ITensor::at(mJointDecodingOutput->logProbs, {batchIdx}); -} - -TensorPtr DecoderState::getSequenceLengths() const -{ - return mJointDecodingOutput->lengths; -} - -TensorPtr DecoderState::getSequenceLengths(SizeType32 batchIdx) const -{ - return ITensor::at(mJointDecodingOutput->lengths, {batchIdx}); -} - -TensorPtr DecoderState::getAllNewTokens() const -{ - return mJointDecodingOutput->newTokensSteps; -} - -TensorPtr DecoderState::getNextDraftTokens() const -{ - return mJointDecodingOutput->speculativeDecodingOutputs->nextDraftTokens; -} - -TensorPtr DecoderState::getPrevDraftTokensLengths() const -{ - return mJointDecodingOutput->speculativeDecodingOutputs->prevDraftTokensLen; -} - -TensorPtr DecoderState::getNextDraftTokensLengths() const -{ - return mJointDecodingOutput->speculativeDecodingOutputs->nextDraftTokensLen; -} - -TensorPtr DecoderState::getAcceptedLengthsCumSum() const -{ - return mJointDecodingOutput->speculativeDecodingOutputs->acceptedLengthsCumSum; -} - -TensorPtr DecoderState::getAcceptedPackedPaths() const -{ - return mJointDecodingOutput->speculativeDecodingOutputs->pathsOffsets; -} - -SizeType32 DecoderState::getMaxNumSequences() const -{ - return mMaxNumSequences; -} - -SizeType32 DecoderState::getMaxBeamWidth() const -{ - return mMaxBeamWidth; -} - -SizeType32 DecoderState::getMaxSequenceLength() const -{ - return mMaxSequenceLength; -} - -SizeType32 DecoderState::getMaxDecodingDecoderTokens() const -{ - return mMaxDecodingDecoderTokens; -} - -SizeType32 DecoderState::getMaxDecodingEngineTokens() const -{ - return mMaxDecodingEngineTokens; -} - -SpeculativeDecodingMode DecoderState::getSpeculativeDecodingMode() const -{ - return mSpeculativeDecodingMode; -} - -ExplicitDraftTokensBuffers::Inputs const& DecoderState::getExplicitDraftTokensBuffers() const -{ - return *mJointDecodingOutput->explicitDraftTokensBuffers; -} - -EagleBuffers::Inputs const& DecoderState::getEagleBuffers() const -{ - return *mJointDecodingOutput->eagleBuffers; -} - -LookaheadDecodingBuffers const& DecoderState::getLookaheadBuffers() const -{ - return *mJointDecodingOutput->lookaheadOutputs; -} - -std::vector const& DecoderState::getNumDecodingEngineTokens() const -{ - return mNumDecodingEngineTokens; -} - -SizeType32 DecoderState::getNumDecodingEngineTokens(SizeType32 batchIdx) const -{ - TLLM_CHECK_WITH_INFO( - batchIdx < mMaxNumSequences, "Batch index %d out of bounds (max %d)", batchIdx, mMaxNumSequences); - return mNumDecodingEngineTokens[batchIdx]; -} - -void DecoderState::setNumDecodingEngineTokens(SizeType32 batchIdx, SizeType32 numTokens) -{ - TLLM_CHECK_WITH_INFO( - batchIdx < mMaxNumSequences, "Batch index %d out of bounds (max %d)", batchIdx, mMaxNumSequences); - mNumDecodingEngineTokens[batchIdx] = numTokens; -} - -BeamSearchBuffers const& DecoderState::getBeamSearchBuffers() const -{ - return *mBeamSearchBuffers; -} - -TensorPtr DecoderState::getCacheIndirectionInput() const -{ - return mJointDecodingInput->cacheIndirection; -} - -TensorPtr DecoderState::getCacheIndirectionOutput() const -{ - return mJointDecodingOutput->cacheIndirection; -} - -std::optional> const& DecoderState::getGenerationSteps() const -{ - return mJointDecodingInput->generationSteps; -} - -void DecoderState::setGenerationSteps(std::vector const& generationSteps) -{ - mJointDecodingInput->generationSteps = generationSteps; -} - -void DecoderState::setBeamWidth(SizeType32 batchIdx, SizeType32 beamWidth) -{ - mJointDecodingInput->beamWidths.at(batchIdx) = beamWidth; -} - -DecodingInput& DecoderState::getJointDecodingInput() const -{ - return *mJointDecodingInput; -} - -DecodingOutput& DecoderState::getJointDecodingOutput() const -{ - return *mJointDecodingOutput; -} - -} // namespace tensorrt_llm::runtime::decoder diff --git a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp deleted file mode 100644 index c5098bf777e0..000000000000 --- a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp +++ /dev/null @@ -1,135 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" -#include "tensorrt_llm/common/tllmDataType.h" - -#include - -tensorrt_llm::runtime::DecodingLayerWorkspace::DecodingLayerWorkspace(std::shared_ptr bufferManager, - tensorrt_llm::layers::DecoderDomain const& decoderDomain, tensorrt_llm::DataType logitsType, - size_t workspaceBufferSizeInBytes) - : mBufferManager(std::move(bufferManager)) - , mBatchSlotsDevice( - mBufferManager->gpu(ITensor::makeShape({decoderDomain.getBatchSize()}), TRTDataType::value)) - , mRuntimeLogitsDevice( - mBufferManager->gpu(ITensor::makeShape({decoderDomain.getBatchSize(), decoderDomain.getMaxDecodingTokens(), - decoderDomain.getBeamWidth(), decoderDomain.getVocabSizePadded()}), - logitsType)) - , mCurandStatesDevice( - mBufferManager->gpu(ITensor::makeShape({decoderDomain.getBatchSize(), sizeof(curandState_t)}))) - , mWorkspaceDeviceBuffer(mBufferManager->gpu(workspaceBufferSizeInBytes)) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("Creating decoding workspace for a maximum batch size of %i, with a scratch space of %lu bytes", - decoderDomain.getBatchSize(), workspaceBufferSizeInBytes); - mBufferManager->getStream().synchronize(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void* tensorrt_llm::runtime::DecodingLayerWorkspace::getRawWorkspaceDevicePtr() const -{ - return mWorkspaceDeviceBuffer->data(); -} - -tensorrt_llm::runtime::DecodingLayerWorkspace::BufferPtr -tensorrt_llm::runtime::DecodingLayerWorkspace::getWorkspaceDeviceBuffer() const -{ - return mWorkspaceDeviceBuffer; -} - -void tensorrt_llm::runtime::DecodingLayerWorkspace::setDeviceBatchSlots(TensorConstPtr const& newBatchSlots) -{ - mBatchSlotsDevice->reshape(newBatchSlots->getShape()); - mBufferManager->copy(*newBatchSlots, *mBatchSlotsDevice); -} - -tensorrt_llm::runtime::SizeType32 const* tensorrt_llm::runtime::DecodingLayerWorkspace::getDeviceBatchSlotsPtr() const -{ - return tensorrt_llm::runtime::bufferCast(*mBatchSlotsDevice); -} - -tensorrt_llm::runtime::DecodingLayerWorkspace::TensorConstPtr -tensorrt_llm::runtime::DecodingLayerWorkspace::getDeviceBatchSlots() const -{ - return mBatchSlotsDevice; -} - -tensorrt_llm::runtime::DecodingLayerWorkspace::TensorPtr -tensorrt_llm::runtime::DecodingLayerWorkspace::getDeviceRuntimeLogits() const -{ - return mRuntimeLogitsDevice; -} - -void tensorrt_llm::runtime::DecodingLayerWorkspace::resize(size_t minSize) -{ - if (mWorkspaceDeviceBuffer->getSizeInBytes() < minSize) - { - mWorkspaceDeviceBuffer->resize(minSize); - } -} - -tensorrt_llm::runtime::DecodingLayerWorkspace::TensorPtr -tensorrt_llm::runtime::DecodingLayerWorkspace::getWorkspaceAsDeviceTensor( - ITensor::Shape shape, tensorrt_llm::DataType type) -{ - auto const sizeInBytes = ITensor::volume(shape) * BufferDataType(type).getSize(); - return std::make_shared>>( - shape, type, BorrowingAllocator{mWorkspaceDeviceBuffer->data(), sizeInBytes}); -} - -void tensorrt_llm::runtime::DecodingLayerWorkspace::initializeDeviceCurandStates( - std::optional> const& randomSeed, tensorrt_llm::runtime::SizeType32 batchSize, - tensorrt_llm::runtime::DecodingLayerWorkspace::TensorConstPtr const& batchSlots, - tensorrt_llm::runtime::DecodingLayerWorkspace::TensorPtr& statesDevice) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - // If runtime argument has single random seed, using this random seed to - // initialize the random table of all sentences. If the argument has - // [batchSize] random seeds, initializing the random table by different - // random seeds respectively. If no random seed, initialize the random table - // of all sentences by 0 directly. - auto const* batchSlotsPtr = tensorrt_llm::runtime::bufferCast(*batchSlots); - auto* curandStateDevicePtr = reinterpret_cast(statesDevice->data()); - if (randomSeed) - { - if (randomSeed->size() == 1) - { - tensorrt_llm::kernels::invokeCurandInitialize( - curandStateDevicePtr, batchSlotsPtr, batchSize, randomSeed->front(), getStream()); - } - else - { - TLLM_CHECK_WITH_INFO(static_cast(randomSeed->size()) == batchSize, - "Random seed vector size mismatch."); - auto randomSeedsDevice = copyToWorkspace(randomSeed.value()); - auto const* randomSeedsDevicePtr = tensorrt_llm::runtime::bufferCast(*randomSeedsDevice); - tensorrt_llm::kernels::invokeCurandBatchInitialize( - curandStateDevicePtr, batchSlotsPtr, batchSize, randomSeedsDevicePtr, getStream()); - } - } - else - { - // Initialize curand states using the default seed 0. - tensorrt_llm::kernels::invokeCurandInitialize(curandStateDevicePtr, batchSlotsPtr, batchSize, 0, getStream()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -cudaStream_t tensorrt_llm::runtime::DecodingLayerWorkspace::getStream() -{ - return mBufferManager->getStream().get(); -} diff --git a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h deleted file mode 100644 index 68d3d54124f5..000000000000 --- a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h +++ /dev/null @@ -1,202 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include - -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/common/workspace.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/tllmBuffers.h" - -namespace tensorrt_llm::runtime -{ - -///@brief A collection of shared resources and data for the decoding layers. -class DecodingLayerWorkspace -{ -public: - using TensorPtr = ITensor::SharedPtr; - using TensorUniquePtr = ITensor::UniquePtr; - using TensorConstPtr = ITensor::SharedConstPtr; - using BufferPtr = IBuffer::SharedPtr; - - DecodingLayerWorkspace(std::shared_ptr bufferManager, layers::DecoderDomain const& decoderDomain, - tensorrt_llm::DataType logitsType, size_t workspaceBufferSizeInBytes); - - DecodingLayerWorkspace() = delete; - - DecodingLayerWorkspace(DecodingLayerWorkspace const& decodingLayerWorkspace) = delete; - - ///@brief Gets a pointer to the start of the shared device workspace. - [[nodiscard]] void* getRawWorkspaceDevicePtr() const; - - ///@brief Gets a pointer to the start of the shared device workspace, as a pointer to the given type. - template - T* getWorkspaceDevicePtrAs() const - { - return reinterpret_cast(mWorkspaceDeviceBuffer->data()); - }; - - ///@brief Gets a pointer to the buffer backing the device workspace. - [[nodiscard]] BufferPtr getWorkspaceDeviceBuffer() const; - - ///@brief Sets the value of the device copy of the batch slots. - void setDeviceBatchSlots(TensorConstPtr const& newBatchSlots); - - ///@brief Gets the pointer to the batch slots on device. - [[nodiscard]] SizeType32 const* getDeviceBatchSlotsPtr() const; - - ///@brief Gets the device tensor containing the batch slots. - [[nodiscard]] TensorConstPtr getDeviceBatchSlots() const; - - ///@brief Gets the device tensor containing the runtime logits. - [[nodiscard]] TensorPtr getDeviceRuntimeLogits() const; - - ///@brief Gets a tensor with the given shape and type at the start of the device workspace. - TensorPtr getWorkspaceAsDeviceTensor(ITensor::Shape shape, tensorrt_llm::DataType type); - - /// @brief A convenience function to copy the content of a standard vector to a device workspace. - template - static void copyToWorkspace(runtime::BufferManager const& bufferManager, std::vector const& src, - runtime::IBuffer::SharedPtr workspace) - { - auto const sizeOfWorkspaceInBytes = workspace->getSizeInBytes(); - auto const sizeOfSrcInBytes = sizeof(T) * src.size(); - TLLM_CHECK_WITH_INFO(sizeOfSrcInBytes <= sizeOfWorkspaceInBytes, - "The size of the workspace (%zu bytes) is insufficient for the data (%zu bytes)", sizeOfWorkspaceInBytes, - sizeOfSrcInBytes); - auto const sizePerWorkspaceElement = BufferDataType(workspace->getDataType()).getSize(); - TLLM_CHECK_WITH_INFO(sizePerWorkspaceElement == 1 || sizePerWorkspaceElement == sizeof(T), - "Copy to typed workspace, but element size mismatched (src: %zu, workspace: %zu)", sizeof(T), - sizePerWorkspaceElement); - runtime::IBuffer::SharedPtr workspaceSlice - = runtime::IBuffer::slice(workspace, 0, sizeOfSrcInBytes / sizePerWorkspaceElement); - bufferManager.copy(src.data(), *workspaceSlice, runtime::MemoryType::kCPU); - } - - /// @brief A convenience function to copy the content of a standard vector to the workspace. - template - TensorPtr copyToWorkspace(std::vector const& src) - { - copyToWorkspace(*mBufferManager, src, mWorkspaceDeviceBuffer); - return getWorkspaceAsDeviceTensor( - ITensor::makeShape({static_cast(src.size())}), TRTDataType::value); - } - - ///@brief Ensures the workspace has at least the provided space in bytes. Does nothing if the workspace is already - /// at least as large. - void resize(size_t minSize); - - ///@brief Given a collection of tuples of tensor shapes and data types, returns the memory aligned size required to - /// contain those tensors. - template - size_t static calculateRequiredWorkspaceSize(Args&&... args) - { - size_t lastTensorOffset = 0; - auto alignedSizeCalculator - = [&lastTensorOffset](std::pair const& tensorDescriptor) - { - auto const& [shape, type] = tensorDescriptor; - auto const sizeInBytes = ITensor::volume(shape) * tensorrt_llm::common::getDTypeSize(type); - auto const sliceEnd = lastTensorOffset + sizeInBytes; - lastTensorOffset = tensorrt_llm::common::alignSize(sliceEnd, tensorrt_llm::common::kCudaMemAlign); - }; - auto argTuple = std::make_tuple(std::forward(args)...); - forEach(alignedSizeCalculator, argTuple); - return lastTensorOffset; - } - - ///@brief Given a collection of tensors, creates tensors with the same shape and data types in the workspace and - /// copies the data from the input tensors to their reflection on device. - template - auto mirrorInWorkspace(Args&&... args) - { - auto* lastTensorEndPtr = reinterpret_cast(mWorkspaceDeviceBuffer->data()); - auto tensorFactory = [&lastTensorEndPtr, this](auto const& tensor) - { - if (tensor == nullptr) - { - return std::unique_ptr>>{}; - } - auto const sizeInBytes = tensor->getSizeInBytes(); - auto const borrowingAllocator = BorrowingAllocator{lastTensorEndPtr, sizeInBytes}; - auto res = std::make_unique>>( - tensor->getShape(), tensor->getDataType(), borrowingAllocator); - auto const sliceEnd = lastTensorEndPtr + sizeInBytes; - lastTensorEndPtr = tensorrt_llm::common::alignPtr(sliceEnd, tensorrt_llm::common::kCudaMemAlign); - mBufferManager->copy(*tensor, *res); - return res; - }; - auto argTuple = std::make_tuple(std::forward(args)...); - - auto res = transform(tensorFactory, argTuple); - std::size_t const numArgs = sizeof...(Args); - std::size_t const sizeInBytes - = lastTensorEndPtr - reinterpret_cast(mWorkspaceDeviceBuffer->data()); - TLLM_LOG_DEBUG("Borrowing %lu bytes of the workspace for %i tensors.", sizeInBytes, numArgs); - return res; - } - - /// @brief A convenience function to initialize curand states from a provided seed. - void initializeDeviceCurandStates(std::optional> const& randomSeed, - runtime::SizeType32 batchSize, TensorConstPtr const& batchSlots, TensorPtr& statesDevice); - -private: - std::shared_ptr mBufferManager; - TensorPtr mBatchSlotsDevice; // - auto static transformImpl(Func&& func, Tuple&& tuple, std::index_sequence) - { - return std::make_tuple(func(std::get(tuple))...); - } - - ///@brief A helper template to apply a function to each element of a tuple and return a tuple of the results. - template - auto static transform(Func&& func, std::tuple const& tuple) - { - return transformImpl(std::forward(func), tuple, std::index_sequence_for{}); - } - - ///@brief A helper template to apply a function to each element of a tuple. - template - void static forEachImpl(Func&& func, Tuple&& tuple, std::index_sequence) - { - (func(std::get(tuple)), ...); - } - - ///@brief A helper template to apply a function to each element of a tuple. - template - void static forEach(Func&& func, std::tuple const& tuple) - { - forEachImpl(std::forward(func), tuple, std::index_sequence_for{}); - } -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/decodingOutput.cpp b/cpp/tensorrt_llm/runtime/decodingOutput.cpp deleted file mode 100644 index 6ff84235dd89..000000000000 --- a/cpp/tensorrt_llm/runtime/decodingOutput.cpp +++ /dev/null @@ -1,75 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/decodingOutput.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -using namespace tensorrt_llm::runtime; - -void DecodingOutput::BeamHypotheses::empty(BufferManager const& manager) -{ - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - auto constexpr nvFloatType = TRTDataType::value; - auto constexpr nvBoolType = TRTDataType::value; - - outputIdsCBA = manager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - logProbsCBA = manager.emptyTensor(MemoryType::kGPU, nvFloatType); - sequenceLengthsCBA = manager.emptyTensor(MemoryType::kGPU, nvSizeType); - cumLogProbsCBA = manager.emptyTensor(MemoryType::kGPU, nvFloatType); - normedScoresCBA = manager.emptyTensor(MemoryType::kGPU, nvFloatType); - numBeamsCBA = manager.emptyTensor(MemoryType::kGPU, nvSizeType); - minNormedScoresCBA = manager.emptyTensor(MemoryType::kGPU, nvFloatType); - batchDones = manager.emptyTensor(MemoryType::kGPU, nvBoolType); -} - -void DecodingOutput::BeamHypotheses::reshape(SizeType32 batchSize, SizeType32 beamWidth, SizeType32 maxSequenceLength) -{ - outputIdsCBA->reshape(ITensor::makeShape({batchSize, 2 * beamWidth, maxSequenceLength})); - logProbsCBA->reshape(ITensor::makeShape({batchSize, 2 * beamWidth, maxSequenceLength})); - sequenceLengthsCBA->reshape(ITensor::makeShape({batchSize, 2 * beamWidth})); - cumLogProbsCBA->reshape(ITensor::makeShape({batchSize, 2 * beamWidth})); - normedScoresCBA->reshape(ITensor::makeShape({batchSize, 2 * beamWidth})); - numBeamsCBA->reshape(ITensor::makeShape({batchSize})); - minNormedScoresCBA->reshape(ITensor::makeShape({batchSize})); - batchDones->reshape(ITensor::makeShape({batchSize})); -} - -void DecodingOutput::BeamHypotheses::init(BufferManager const& manager, TokenIdType endId) -{ - kernels::invokeFill(*outputIdsCBA, endId, manager.getStream()); - manager.setZero(*logProbsCBA); - manager.setZero(*sequenceLengthsCBA); - manager.setZero(*cumLogProbsCBA); - manager.setZero(*normedScoresCBA); - manager.setZero(*numBeamsCBA); - manager.setZero(*minNormedScoresCBA); - manager.setZero(*batchDones); -} - -DecodingOutput::BeamHypotheses DecodingOutput::BeamHypotheses::slice(SizeType32 batchIndex, SizeType32 size) const -{ - DecodingOutput::BeamHypotheses bh{}; - bh.outputIdsCBA = ITensor::slice(outputIdsCBA, batchIndex, size); - bh.logProbsCBA = ITensor::slice(logProbsCBA, batchIndex, size); - bh.sequenceLengthsCBA = ITensor::slice(sequenceLengthsCBA, batchIndex, size); - bh.cumLogProbsCBA = ITensor::slice(cumLogProbsCBA, batchIndex, size); - bh.normedScoresCBA = ITensor::slice(normedScoresCBA, batchIndex, size); - bh.numBeamsCBA = ITensor::slice(numBeamsCBA, batchIndex, size); - bh.minNormedScoresCBA = ITensor::slice(minNormedScoresCBA, batchIndex, size); - bh.batchDones = ITensor::slice(batchDones, batchIndex, size); - return bh; -} diff --git a/cpp/tensorrt_llm/runtime/eagleBuffers.cpp b/cpp/tensorrt_llm/runtime/eagleBuffers.cpp deleted file mode 100644 index e0f2198c3e58..000000000000 --- a/cpp/tensorrt_llm/runtime/eagleBuffers.cpp +++ /dev/null @@ -1,602 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/eagleBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -namespace tksd = tensorrt_llm::kernels::speculative_decoding; - -namespace tensorrt_llm::runtime -{ - -void EagleBuffers::Inputs::create(SizeType32 maxNumSequences, BufferManager const& manager, - ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - auto const& speculativeDecodingModule = modelConfig.getSpeculativeDecodingModule(); - auto const maxNumPaths = speculativeDecodingModule.getMaxNumPaths(); - auto const maxPathLen = speculativeDecodingModule.getMaxPathLen(); - auto const maxDecodingTokens = speculativeDecodingModule.getMaxDecodingTokens(); - auto const maxDecodingDraftTokens = speculativeDecodingModule.getMaxDecodingDraftTokens(); - auto const numEagleLayers = speculativeDecodingModule.getMaxDraftPathLen(); - auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - - temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kFLOAT); - randomDataSample = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kFLOAT); - randomDataValidation - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens}), tensorrt_llm::DataType::kFLOAT); - draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - draftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); - draftPathsHost = BufferManager::pinnedPool( - ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); - specDecodingGenerationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - specDecodingGenerationLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - specDecodingPackedMasks - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)}), - tensorrt_llm::DataType::kINT32); - specDecodingPositionOffsets - = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); - - eagleNetCtxRequestTypesHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - eagleNetCtxContextLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - eagleNetCtxPastKeyValueLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - eagleNetGenRequestTypesHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - eagleNetGenContextLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - eagleNetGenPastKeyValueLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - inputGenTokensHost = BufferManager::pinnedPool( - ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); - chunkedContextNextTokens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - // Eagle-2 - useDynamicTreeHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - dynamicTreeMaxTopKHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - prevScores - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), tensorrt_llm::DataType::kFLOAT); - currentExpandIndices = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - allLayersScores = manager.gpu( - ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - tensorrt_llm::DataType::kFLOAT); - allLayersDraftTokenIds = manager.gpu( - ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - TRTTokenIdType); - allLayersDraftTokenIdsPredecessor = manager.gpu( - ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - TRTTokenIdType); -} - -EagleBuffers::EagleBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK_WITH_INFO(maxBeamWidth == 1, "EAGLE does not support beam search"); - - auto const maxNumSequences = maxBatchSize; - - auto const eagleModule = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const numPaths = eagleModule->getMaxNumPaths(); - auto const pathLen = eagleModule->getMaxPathLen(); - auto const maxDecodingDraftTokens = eagleModule->getMaxDecodingDraftTokens(); - auto const numEagleLayers = eagleModule->getMaxDraftPathLen(); - auto const maxNonLeafNodesPerLayer = eagleModule->getMaxNonLeafNodesPerLayer(); - - auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - - // input tensors - engineInputs.temperatures = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); - engineInputs.posteriorAlpha = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); - engineInputs.posteriorThreshold = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); - posteriorAlphaHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kFLOAT); - posteriorThresholdHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kFLOAT); - greedySamplingHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - engineInputs.draftTokens - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - engineInputs.draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - engineInputs.draftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), tensorrt_llm::DataType::kINT32); - - engineInputs.specDecodingGenerationLengths - = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.specDecodingPositionOffsets - = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.specDecodingPackedMasks - = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - - engineInputs.randomDataSample = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); - engineInputs.randomDataValidation = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); - - engineInputs.eagleNetCtxRequestTypesHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.eagleNetCtxContextLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.eagleNetCtxPastKeyValueLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.eagleNetGenRequestTypesHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.eagleNetGenContextLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.eagleNetGenPastKeyValueLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.inputGenTokensHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - engineInputs.chunkedContextNextTokens - = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.useSpecDecoding = BufferManager::cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - bufferCast(*engineInputs.useSpecDecoding)[0] = 1; - chunkedContextNextTokensHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); - - // Eagle-2 - engineInputs.useDynamicTreeHost - = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - engineInputs.dynamicTreeMaxTopKHost - = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - engineInputs.prevScores - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), tensorrt_llm::DataType::kFLOAT); - engineInputs.currentExpandIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - engineInputs.allLayersScores = manager.gpu( - ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - tensorrt_llm::DataType::kFLOAT); - engineInputs.allLayersDraftTokenIds = manager.gpu( - ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - TRTTokenIdType); - engineInputs.allLayersDraftTokenIdsPredecessor = manager.gpu( - ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - TRTTokenIdType); - - // output tensors - engineOutputs.nextDraftTokens - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), TRTTokenIdType); - engineOutputs.nextDraftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - engineOutputs.nextDraftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), tensorrt_llm::DataType::kINT32); - - engineOutputs.acceptedTokens - = manager.gpu(ITensor::makeShape({maxNumSequences, pathLen}), tensorrt_llm::DataType::kINT32); - engineOutputs.acceptedLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - engineOutputs.acceptedPaths = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - engineOutputs.chunkedContextNextTokens - = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - - // helper tensors - scanReduceTempStorageBytes = tksd::invokeScanReduceGenerationLengths( - maxNumSequences, nullptr, nullptr, 0, nullptr, nullptr, manager.getStream().get()); - scanReduceTempStorage = manager.gpu(scanReduceTempStorageBytes); - - cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - maxGenerationLength = manager.gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - // pre-allocate empty tensors - reshape(0, maxNumSequences, modelConfig); - - // Init defaults - auto const defaultConfig = decodingConfig.getEagleConfig().value_or(tensorrt_llm::executor::EagleConfig()); - mDoGreedySampling = defaultConfig.isGreedySampling(); - mDefaultPosteriorThreshold = defaultConfig.getPosteriorThreshold().value_or(mDefaultPosteriorThreshold); - bufferCast(*greedySamplingHost)[0] = mDoGreedySampling; - - auto const useDynamicTree = defaultConfig.useDynamicTree(); - auto const dynamicTreeMaxTopK = defaultConfig.getDynamicTreeMaxTopK().value_or(-1); - - if (useDynamicTree) - { - TLLM_LOG_WARNING("EAGLE-2 is still under the experimental stage."); - TLLM_CHECK_WITH_INFO(dynamicTreeMaxTopK > 0, - "When using Eagle-2, dynamicTreeMaxTopK should greater than 0. Now dynamicTreeMaxTopK is %d", - dynamicTreeMaxTopK); - TLLM_CHECK_WITH_INFO(maxNonLeafNodesPerLayer >= dynamicTreeMaxTopK, - "When using Eagle-2, maxNonLeafNodesPerLayer should be greater or equal to dynamicTreeMaxTopK. Now " - "maxNonLeafNodesPerLayer is %d, and dynamicTreeMaxTopK is %d", - maxNonLeafNodesPerLayer, dynamicTreeMaxTopK); - TLLM_CHECK_WITH_INFO(maxDecodingDraftTokens >= dynamicTreeMaxTopK, - "When using Eagle-2, maxDecodingDraftTokens should be greater or equal to dynamicTreeMaxTopK. Now " - "maxDecodingDraftTokens is %d, and dynamicTreeMaxTopK is %d", - maxDecodingDraftTokens, dynamicTreeMaxTopK); - } - - // Eagle-2 config - bufferCast(*engineInputs.useDynamicTreeHost)[0] = SizeType32(useDynamicTree); - bufferCast(*engineInputs.dynamicTreeMaxTopKHost)[0] = dynamicTreeMaxTopK; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EagleBuffers::reshape( - SizeType32 numCtxSequences, SizeType32 numGenSequences, runtime::ModelConfig const& modelConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const numSequences = numCtxSequences + numGenSequences; - - auto const eagleModule = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const maxDecodingTokens = eagleModule->getMaxDecodingTokens(); - - // input tensors - engineInputs.temperatures->reshape(ITensor::makeShape({numSequences})); - engineInputs.posteriorAlpha->reshape(ITensor::makeShape({numSequences})); - engineInputs.posteriorThreshold->reshape(ITensor::makeShape({numSequences})); - posteriorAlphaHost->reshape(ITensor::makeShape({numSequences})); - posteriorThresholdHost->reshape(ITensor::makeShape({numSequences})); - - auto draftTokensShape = engineInputs.draftTokens->getShape(); - draftTokensShape.d[0] = numSequences; - engineInputs.draftTokens->reshape(draftTokensShape); - auto draftLensShape = engineInputs.draftLens->getShape(); - draftLensShape.d[0] = numSequences; - engineInputs.draftLens->reshape(draftLensShape); - auto draftPathsShape = engineInputs.draftPaths->getShape(); - draftPathsShape.d[0] = numSequences; - engineInputs.draftPaths->reshape(draftPathsShape); - - engineInputs.specDecodingGenerationLengths->reshape(ITensor::makeShape({numGenSequences})); - engineInputs.specDecodingPositionOffsets->reshape(ITensor::makeShape({numGenSequences, maxDecodingTokens})); - engineInputs.specDecodingPackedMasks->reshape( - ITensor::makeShape({numGenSequences * maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)})); - - engineInputs.randomDataSample->reshape(ITensor::makeShape({numSequences})); - engineInputs.randomDataValidation->reshape(ITensor::makeShape({numSequences, maxDecodingTokens})); - - engineInputs.eagleNetCtxRequestTypesHost->reshape(ITensor::makeShape({numSequences})); - engineInputs.eagleNetCtxContextLengthsHost->reshape(ITensor::makeShape({numSequences})); - engineInputs.eagleNetCtxPastKeyValueLengthsHost->reshape(ITensor::makeShape({numSequences})); - engineInputs.eagleNetGenRequestTypesHost->reshape(ITensor::makeShape({numSequences})); - engineInputs.eagleNetGenContextLengthsHost->reshape(ITensor::makeShape({numSequences})); - engineInputs.eagleNetGenPastKeyValueLengthsHost->reshape(ITensor::makeShape({numSequences})); - engineInputs.inputGenTokensHost->reshape(ITensor::makeShape({numSequences * maxDecodingTokens})); - engineInputs.chunkedContextNextTokens->reshape(ITensor::makeShape({numSequences})); - // Eagle-2 - // Reshape prevScores - auto prevScoresShape = engineInputs.prevScores->getShape(); - prevScoresShape.d[0] = numSequences; - engineInputs.prevScores->reshape(prevScoresShape); - // Reshape currentExpandIndices - auto currentExpandIndicesShape = engineInputs.currentExpandIndices->getShape(); - currentExpandIndicesShape.d[0] = numSequences; - engineInputs.currentExpandIndices->reshape(currentExpandIndicesShape); - // Reshape allLayersScores - auto allLayersScoresShape = engineInputs.allLayersScores->getShape(); - allLayersScoresShape.d[0] = numSequences; - engineInputs.allLayersScores->reshape(allLayersScoresShape); - // Reshape allLayersDraftTokenIds - auto allLayersDraftTokenIdsShape = engineInputs.allLayersDraftTokenIds->getShape(); - allLayersDraftTokenIdsShape.d[0] = numSequences; - engineInputs.allLayersDraftTokenIds->reshape(allLayersDraftTokenIdsShape); - // Reshape allLayersDraftTokenIdsPredecessor - auto allLayersDraftTokenIdsPredecessorShape = engineInputs.allLayersDraftTokenIdsPredecessor->getShape(); - allLayersDraftTokenIdsPredecessorShape.d[0] = numSequences; - engineInputs.allLayersDraftTokenIdsPredecessor->reshape(allLayersDraftTokenIdsPredecessorShape); - - chunkedContextNextTokensHost->reshape(ITensor::makeShape({numSequences})); - engineOutputs.chunkedContextNextTokens->reshape(ITensor::makeShape({numSequences})); - - cumSumGenerationLengths->reshape(ITensor::makeShape({numSequences + 1})); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleBuffers::setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, - SizeType32 vocabSizePadded, ITensor const& seqSlots, EagleBuffers::Inputs const& draftBuffers, - runtime::EagleModule const& eagleModule, runtime::BufferManager const& manager) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - using runtime::bufferCast; - - auto const numCtxSequences = static_cast(contextRequests.size()); - auto const numGenSequences = static_cast(genRequests.size()); - - tksd::PackEagleParams params; - params.batchSize = numCtxSequences + numGenSequences; - params.maxNumPaths = eagleModule.getMaxNumPaths(); - params.maxDecodingTokens = eagleModule.getMaxDecodingTokens(); - params.maxPathLength = eagleModule.getMaxPathLen(); - params.numContextRequests = numCtxSequences; - params.numGenerationRequests = numGenSequences; - - params.batchSlots = bufferCast(seqSlots); - - // Outputs from decoder -- inputs to the packing kernel - params.inputTemperatures = bufferCast(*draftBuffers.temperatures); - params.inputRandomDataSample = bufferCast(*draftBuffers.randomDataSample); - params.inputRandomDataValidation = bufferCast(*draftBuffers.randomDataValidation); - - params.inputNextDraftTokens = bufferCast(*draftBuffers.draftTokens); - params.inputNextDraftPaths = bufferCast(*draftBuffers.draftPaths); - - params.inputSpecDecodingGenerationLengths = bufferCast(*draftBuffers.specDecodingGenerationLengths); - params.inputSpecDecodingPositionOffsets = bufferCast(*draftBuffers.specDecodingPositionOffsets); - params.inputSpecDecodingPackedMasks = bufferCast(*draftBuffers.specDecodingPackedMasks); - - // Outputs of the packing kernel -- inputs to the engine - params.outputTemperatures = bufferCast(*engineInputs.temperatures); - params.outputRandomDataSample = bufferCast(*engineInputs.randomDataSample); - params.outputRandomDataValidation = bufferCast(*engineInputs.randomDataValidation); - - params.outputNextDraftTokens = bufferCast(*engineInputs.draftTokens); - params.outputNextDraftLens = bufferCast(*engineInputs.draftLens); - params.outputNextDraftPaths = bufferCast(*engineInputs.draftPaths); - - params.outputSpecDecodingGenerationLengths = bufferCast(*engineInputs.specDecodingGenerationLengths); - params.outputSpecDecodingPositionOffsets = bufferCast(*engineInputs.specDecodingPositionOffsets); - params.outputSpecDecodingPackedMasks = bufferCast(*engineInputs.specDecodingPackedMasks); - - params.maxGenerationLength = bufferCast(*maxGenerationLength); - params.cumSumGenerationLengths = bufferCast(*cumSumGenerationLengths); - - params.checkParams(); - - // Pack tensors from batch slot position to continuous array - tksd::invokePackEagleGenerationLengths(params, manager.getStream().get()); - - if (numGenSequences) - { - // Compute inclusive sum and max - tksd::invokeScanReduceGenerationLengths(numGenSequences, - bufferCast(*engineInputs.specDecodingGenerationLengths), - bufferCast(*scanReduceTempStorage), scanReduceTempStorageBytes, - bufferCast(*cumSumGenerationLengths), bufferCast(*maxGenerationLength), - manager.getStream().get()); - } - - // Pack tensors from batch slot position to continuous array - tksd::invokePackEagle(params, manager.getStream().get()); - - // Pack host data. - SizeType32 maxGenerationLengthHostValue{-1}; - SizeType32 numGenerationTokens{0}; - SizeType32 batchIdx{0}; - - auto chunkedContextNextTokensHostPtr = bufferCast(*chunkedContextNextTokensHost); - std::fill(chunkedContextNextTokensHostPtr, chunkedContextNextTokensHostPtr + params.batchSize, -1); - - auto setupEagleNetHostBuffers = [this, &draftBuffers](SizeType32 batchIdx, SizeType32 batchSlot) - { - bufferCast(*this->engineInputs.eagleNetCtxRequestTypesHost)[batchIdx] - = bufferCast(*draftBuffers.eagleNetCtxRequestTypesHost)[batchSlot]; - - bufferCast(*this->engineInputs.eagleNetCtxContextLengthsHost)[batchIdx] - = bufferCast(*draftBuffers.eagleNetCtxContextLengthsHost)[batchSlot]; - - bufferCast(*this->engineInputs.eagleNetCtxPastKeyValueLengthsHost)[batchIdx] - = bufferCast(*draftBuffers.eagleNetCtxPastKeyValueLengthsHost)[batchSlot]; - - bufferCast(*this->engineInputs.eagleNetGenRequestTypesHost)[batchIdx] - = bufferCast(*draftBuffers.eagleNetGenRequestTypesHost)[batchSlot]; - - bufferCast(*this->engineInputs.eagleNetGenContextLengthsHost)[batchIdx] - = bufferCast(*draftBuffers.eagleNetGenContextLengthsHost)[batchSlot]; - - bufferCast(*this->engineInputs.eagleNetGenPastKeyValueLengthsHost)[batchIdx] - = bufferCast(*draftBuffers.eagleNetGenPastKeyValueLengthsHost)[batchSlot]; - }; - - auto posteriorAlphaHostPtr = bufferCast(*posteriorAlphaHost); - auto posteriorThresholdHostPtr = bufferCast(*posteriorThresholdHost); - auto setPosteriorThresholds - = [this, posteriorAlphaHostPtr, posteriorThresholdHostPtr](LlmRequestPtr const& llmReq, SizeType32 batchIdx) - { - auto const eagleConfig = llmReq->getEagleConfig(); - - float posteriorThreshold{this->mDefaultPosteriorThreshold}; - if (eagleConfig.has_value()) - { - posteriorThreshold = eagleConfig->getPosteriorThreshold().value_or(posteriorThreshold); - } - posteriorAlphaHostPtr[batchIdx] = std::sqrt(posteriorThreshold); - posteriorThresholdHostPtr[batchIdx] = posteriorThreshold; - }; - - for (auto const& llmReq : contextRequests) - { - if (llmReq->isLastContextChunk()) - { - auto const batchSlot = params.batchSlots[batchIdx]; - setupEagleNetHostBuffers(batchIdx, batchSlot); - - auto draftTokens = ITensor::slice(engineInputs.draftTokens, batchIdx, 1); - runtime::kernels::invokeFill(*draftTokens, -1, manager.getStream()); - } - else - { - auto const contextChunkSize = llmReq->getContextChunkSize(); - auto const beginCompute = llmReq->getContextCurrentPosition(); - auto const endCompute = beginCompute + contextChunkSize; - - // Fill values for requests with chunked context as their decoder setup step is skipped. - bufferCast(*engineInputs.eagleNetCtxRequestTypesHost)[batchIdx] = 0; - bufferCast(*engineInputs.eagleNetCtxContextLengthsHost)[batchIdx] = contextChunkSize; - bufferCast(*engineInputs.eagleNetCtxPastKeyValueLengthsHost)[batchIdx] - = beginCompute + contextChunkSize; - - bufferCast(*engineInputs.eagleNetGenRequestTypesHost)[batchIdx] = 1; - bufferCast(*engineInputs.eagleNetGenContextLengthsHost)[batchIdx] - = beginCompute + contextChunkSize; - bufferCast(*engineInputs.eagleNetGenPastKeyValueLengthsHost)[batchIdx] - = beginCompute + contextChunkSize; - - // Setup fake path - TensorPtr draftPathsHostSlice = ITensor::at(engineInputs.draftPathsHost, {batchIdx, 1}); - - for (SizeType32 ti = 0; ti < eagleModule.getMaxPathLen(); ++ti) - { - bufferCast(*draftPathsHostSlice)[ti] = ti; - } - - TensorPtr draftPathsBatchSlice = ITensor::slice(engineInputs.draftPaths, batchIdx, 1); - draftPathsBatchSlice->squeeze(0); - kernels::invokeFill(*draftPathsBatchSlice, -1, manager.getStream()); - TensorPtr draftPathsBatchPathSlice = ITensor::slice(draftPathsBatchSlice, 0, 1); - manager.copy(*draftPathsHostSlice, *draftPathsBatchPathSlice); - - auto const& reqTokens = llmReq->getTokens(0); - chunkedContextNextTokensHostPtr[batchIdx] = reqTokens[endCompute]; - } - - setPosteriorThresholds(llmReq, batchIdx); - - ++batchIdx; - } - - for (auto const& llmReq : genRequests) - { - auto const batchSlot = params.batchSlots[batchIdx]; - setupEagleNetHostBuffers(batchIdx, batchSlot); - setPosteriorThresholds(llmReq, batchIdx); - - auto const generationLength - = bufferCast(*draftBuffers.specDecodingGenerationLengthsHost)[batchSlot]; - maxGenerationLengthHostValue = std::max(maxGenerationLengthHostValue, generationLength); - numGenerationTokens += generationLength; - - ++batchIdx; - } - - if (maxGenerationLengthHostValue <= 0) - { - maxGenerationLengthHostValue = params.maxDecodingTokens; - } - - auto specDecodingPositionOffsetsShape = engineInputs.specDecodingPositionOffsets->getShape(); - specDecodingPositionOffsetsShape.d[1] = maxGenerationLengthHostValue; - engineInputs.specDecodingPositionOffsets->reshape(specDecodingPositionOffsetsShape); - - auto inputGenTokensHostShape = engineInputs.inputGenTokensHost->getShape(); - inputGenTokensHostShape.d[0] = numGenerationTokens; - engineInputs.inputGenTokensHost->reshape(inputGenTokensHostShape); - - manager.copy(*chunkedContextNextTokensHost, *engineInputs.chunkedContextNextTokens); - manager.copy(*chunkedContextNextTokensHost, *engineOutputs.chunkedContextNextTokens); - manager.copy(*posteriorAlphaHost, *engineInputs.posteriorAlpha); - manager.copy(*posteriorThresholdHost, *engineInputs.posteriorThreshold); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EagleBuffers::setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, - ITensor const& requestTypes, ITensor const& seqSlots, EagleBuffers::Inputs const& draftBuffers, - BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const eagleModule - = std::dynamic_pointer_cast(modelConfig.getSpeculativeDecodingModulePtr()); - - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const dtype = modelConfig.getDataType(); - - switch (dtype) - { - case tensorrt_llm::DataType::kFLOAT: - setFromInputs( - contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); - break; - case tensorrt_llm::DataType::kHALF: - setFromInputs( - contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); - break; - case tensorrt_llm::DataType::kBF16: - setFromInputs<__nv_bfloat16>( - contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); - break; - default: TLLM_THROW("DataType %d not supported in EagleBuffers", static_cast(dtype)); break; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EagleBuffers::insertInputTensors( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::WorldConfig const& /* worldConfig */) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - // inputs - inputBuffers.insert_or_assign("greedy_sampling", greedySamplingHost); - inputBuffers.insert_or_assign("eagle_temperature", engineInputs.temperatures); - inputBuffers.insert_or_assign("posterior_alpha", engineInputs.posteriorAlpha); - inputBuffers.insert_or_assign("posterior_threshold", engineInputs.posteriorThreshold); - - inputBuffers.insert_or_assign("spec_decoding_generation_lengths", engineInputs.specDecodingGenerationLengths); - inputBuffers.insert_or_assign("spec_decoding_position_offsets", engineInputs.specDecodingPositionOffsets); - inputBuffers.insert_or_assign("spec_decoding_packed_mask", engineInputs.specDecodingPackedMasks); - - inputBuffers.insert_or_assign("rand_data_sample", engineInputs.randomDataSample); - inputBuffers.insert_or_assign("rand_data_validation", engineInputs.randomDataValidation); - - inputBuffers.insert_or_assign("draft_tokens", engineInputs.draftTokens); - inputBuffers.insert_or_assign("draft_lens", engineInputs.draftLens); - inputBuffers.insert_or_assign("draft_paths", engineInputs.draftPaths); - - inputBuffers.insert_or_assign("host_ctx_eagle_net_request_types", engineInputs.eagleNetCtxRequestTypesHost); - inputBuffers.insert_or_assign("host_ctx_eagle_net_context_lengths", engineInputs.eagleNetCtxContextLengthsHost); - inputBuffers.insert_or_assign( - "host_ctx_eagle_net_past_key_value_lengths", engineInputs.eagleNetCtxPastKeyValueLengthsHost); - inputBuffers.insert_or_assign("host_gen_eagle_net_request_types", engineInputs.eagleNetGenRequestTypesHost); - inputBuffers.insert_or_assign("host_gen_eagle_net_context_lengths", engineInputs.eagleNetGenContextLengthsHost); - inputBuffers.insert_or_assign( - "host_gen_eagle_net_past_key_value_lengths", engineInputs.eagleNetGenPastKeyValueLengthsHost); - inputBuffers.insert_or_assign("input_gen_tokens", engineInputs.inputGenTokensHost); - inputBuffers.insert_or_assign("chunked_context_next_tokens", engineInputs.chunkedContextNextTokens); - // For Eagle-2 - inputBuffers.insert_or_assign("use_dynamic_tree", engineInputs.useDynamicTreeHost); - inputBuffers.insert_or_assign("spec_decoding_use", engineInputs.useSpecDecoding); - inputBuffers.insert_or_assign("dynamic_tree_max_topK", engineInputs.dynamicTreeMaxTopKHost); - inputBuffers.insert_or_assign("prev_scores", engineInputs.prevScores); - inputBuffers.insert_or_assign("current_expand_indices", engineInputs.currentExpandIndices); - inputBuffers.insert_or_assign("all_layers_scores", engineInputs.allLayersScores); - inputBuffers.insert_or_assign("all_layers_draft_token_ids", engineInputs.allLayersDraftTokenIds); - inputBuffers.insert_or_assign( - "all_layers_draft_token_ids_predecessor", engineInputs.allLayersDraftTokenIdsPredecessor); - - // outputs - outputBuffers.insert_or_assign("next_draft_tokens", engineOutputs.nextDraftTokens); - outputBuffers.insert_or_assign("next_draft_lens", engineOutputs.nextDraftLens); - outputBuffers.insert_or_assign("next_draft_paths", engineOutputs.nextDraftPaths); - - outputBuffers.insert_or_assign("accepted_tokens", engineOutputs.acceptedTokens); - outputBuffers.insert_or_assign("num_accepted_tokens", engineOutputs.acceptedLens); - outputBuffers.insert_or_assign("accepted_paths", engineOutputs.acceptedPaths); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp b/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp deleted file mode 100644 index 89c74e6f9349..000000000000 --- a/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp +++ /dev/null @@ -1,382 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/explicitDraftTokensBuffers.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -namespace tksd = tensorrt_llm::kernels::speculative_decoding; - -namespace tensorrt_llm::runtime -{ - -void ExplicitDraftTokensBuffers::Inputs::create(SizeType32 maxNumSequences, BufferManager const& manager, - ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - auto const& speculativeDecodingModule = modelConfig.getSpeculativeDecodingModule(); - auto const maxNumPaths = speculativeDecodingModule.getMaxNumPaths(); - auto const maxDraftPathLen = speculativeDecodingModule.getMaxDraftPathLen(); - auto const maxPathLen = speculativeDecodingModule.getMaxPathLen(); - auto const maxDecodingTokens = speculativeDecodingModule.getMaxDecodingTokens(); - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - auto const dtype = modelConfig.getDataType(); - - maxGenLengthHost = manager.pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), dtype); - positionIdsBase = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - generationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - generationLengthsHost = manager.pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); - randomDataSample = manager.gpu(ITensor::makeShape({maxNumSequences}), dtype); - randomDataValidation = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxDraftPathLen}), dtype); - draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), TRTTokenIdType); - draftIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); - draftProbs - = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxDraftPathLen, vocabSizePadded}), dtype); - packedMasks - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)}), - tensorrt_llm::DataType::kINT32); - positionIds - = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); -} - -ExplicitDraftTokensBuffers::ExplicitDraftTokensBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK_WITH_INFO(maxBeamWidth == 1, "Explicit draft tokens does not support beam search"); - - auto const maxNumSequences = maxBatchSize; - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const explicitDraftTokensModule - = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const numBeams = explicitDraftTokensModule->getMaxNumPaths(); - auto const beamDraftLength = explicitDraftTokensModule->getMaxDraftPathLen(); - auto const beamLength = explicitDraftTokensModule->getMaxPathLen(); // beamDraftLength + 1 - - auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - auto const dtype = modelConfig.getDataType(); - - // input tensors - engineInputs.requestTypesDevice = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.temperatures = manager.emptyTensor(runtime::MemoryType::kGPU, dtype); - - engineInputs.draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), TRTTokenIdType); - engineInputs.draftIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), tensorrt_llm::DataType::kINT32); - engineInputs.draftProbs - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamDraftLength, vocabSizePadded}), dtype); - - engineInputs.generationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.positionIds = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.positionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.packedMasks = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - - engineInputs.randomDataSample = manager.emptyTensor(runtime::MemoryType::kGPU, dtype); - engineInputs.randomDataValidation = manager.emptyTensor(runtime::MemoryType::kGPU, dtype); - engineInputs.positionIdsBase = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.useSpecDecoding = manager.cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - bufferCast(*engineInputs.useSpecDecoding)[0] = 1; - - // output tensors - engineOutputs.nextDraftTokens - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), TRTTokenIdType); - engineOutputs.nextDraftIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), tensorrt_llm::DataType::kINT32); - engineOutputs.nextDraftProbs - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamDraftLength, vocabSizePadded}), dtype); - - engineOutputs.maxGenToken = manager.gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - engineOutputs.totalGenToken = manager.gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - engineOutputs.nextGenerationLengths - = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineOutputs.nextPositionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineOutputs.masks = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kBOOL); - - engineOutputs.nextFlatTokens = manager.emptyTensor(runtime::MemoryType::kGPU, TRTTokenIdType); - engineOutputs.bestPathLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineOutputs.bestPathIndices = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineOutputs.packedPositionIds = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - - // helper tensors - auto const& stream = manager.getStream(); - scanTempStorageBytes - = tksd::invokeScanGenerationLengths(nullptr, 0, nullptr, nullptr, maxNumSequences, stream.get()); - scanTempStorage = manager.gpu(scanTempStorageBytes); - cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - - // pre-allocate empty tensors - reshape(0, maxNumSequences, modelConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void ExplicitDraftTokensBuffers::reshape( - SizeType32 numCtxSequences, SizeType32 numGenSequences, runtime::ModelConfig const& modelConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const numSequences = numCtxSequences + numGenSequences; - - auto const explicitDraftTokensModule - = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const numBeams = explicitDraftTokensModule->getMaxNumPaths(); - auto const beamDraftLength = explicitDraftTokensModule->getMaxDraftPathLen(); - auto const maxDecodingTokens = explicitDraftTokensModule->getMaxDecodingTokens(); - - // input tensors - engineInputs.requestTypesDevice->reshape(ITensor::makeShape({numSequences})); - engineInputs.temperatures->reshape(ITensor::makeShape({numSequences})); - - auto draftTokensShape = engineInputs.draftTokens->getShape(); - draftTokensShape.d[0] = numGenSequences; - engineInputs.draftTokens->reshape(draftTokensShape); - auto draftIndicesShape = engineInputs.draftIndices->getShape(); - draftIndicesShape.d[0] = numGenSequences; - engineInputs.draftIndices->reshape(draftIndicesShape); - auto draftProbsShape = engineInputs.draftProbs->getShape(); - draftProbsShape.d[0] = numGenSequences; - engineInputs.draftProbs->reshape(draftProbsShape); - - engineInputs.generationLengths->reshape(ITensor::makeShape({numGenSequences})); - engineInputs.positionIds->reshape(ITensor::makeShape({numSequences * maxDecodingTokens})); - engineInputs.positionOffsets->reshape(ITensor::makeShape({numGenSequences, maxDecodingTokens})); - engineInputs.packedMasks->reshape( - ITensor::makeShape({numGenSequences * maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)})); - - engineInputs.randomDataSample->reshape(ITensor::makeShape({numSequences})); - engineInputs.randomDataValidation->reshape(ITensor::makeShape({numGenSequences, numBeams, beamDraftLength})); - engineInputs.positionIdsBase->reshape(ITensor::makeShape({numSequences})); - - // output tensors - engineOutputs.nextGenerationLengths->reshape(ITensor::makeShape({numSequences})); - engineOutputs.nextPositionOffsets->reshape(ITensor::makeShape({numSequences, maxDecodingTokens})); - engineOutputs.masks->reshape(ITensor::makeShape({numSequences, maxDecodingTokens, maxDecodingTokens})); - - engineOutputs.nextFlatTokens->reshape(ITensor::makeShape({numSequences * maxDecodingTokens})); - engineOutputs.bestPathLengths->reshape(ITensor::makeShape({numSequences})); - engineOutputs.bestPathIndices->reshape(ITensor::makeShape({numSequences})); - engineOutputs.packedPositionIds->reshape(ITensor::makeShape({numSequences * maxDecodingTokens})); - - cumSumGenerationLengths->reshape(ITensor::makeShape({numSequences})); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void ExplicitDraftTokensBuffers::setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, - SizeType32 vocabSizePadded, ITensor const& seqSlots, ExplicitDraftTokensBuffers::Inputs const& draftBuffers, - ITensor const& contextPositionIds, runtime::ExplicitDraftTokensModule const& explicitDraftTokensModule, - runtime::CudaStream const& stream) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - using runtime::bufferCast; - - tksd::PackExplicitDraftTokensParams params; - params.batchSize = numCtxSequences + numGenSequences; - params.numPaths = explicitDraftTokensModule.getMaxNumPaths(); - params.maxPathLength = explicitDraftTokensModule.getMaxPathLen(); - params.vocabSize = vocabSizePadded; - params.numContextRequests = numCtxSequences; - params.numGenerationRequests = numGenSequences; - params.numContextTokens = contextPositionIds.getShape().d[0]; - - params.batchSlots = bufferCast(seqSlots); - - params.maxGenerationLength = bufferCast(*engineOutputs.maxGenToken); - - params.inputTemperatures = bufferCast(*draftBuffers.temperatures); - params.inputPositionIdsBase = bufferCast(*draftBuffers.positionIdsBase); - params.inputGenerationLengths = bufferCast(*draftBuffers.generationLengths); - params.inputRandomDataSample = bufferCast(*draftBuffers.randomDataSample); - params.inputRandomDataValidation = bufferCast(*draftBuffers.randomDataValidation); - params.inputNextDraftTokens = bufferCast(*draftBuffers.draftTokens); - params.inputNextDraftIndices = bufferCast(*draftBuffers.draftIndices); - params.inputDraftProbs = bufferCast(*draftBuffers.draftProbs); - params.inputPackedMask = bufferCast(*draftBuffers.packedMasks); - params.inputPositionIds = bufferCast(*draftBuffers.positionIds); - - params.outputTemperatures = bufferCast(*engineInputs.temperatures); - params.outputPositionIdsBase = bufferCast(*engineInputs.positionIdsBase); - params.outputGenerationLengths = bufferCast(*engineInputs.generationLengths); - params.outputRandomDataSample = bufferCast(*engineInputs.randomDataSample); - params.outputRandomDataValidation = bufferCast(*engineInputs.randomDataValidation); - params.outputNextDraftTokens = bufferCast(*engineInputs.draftTokens); - params.outputNextDraftIndices = bufferCast(*engineInputs.draftIndices); - params.outputDraftProbs = bufferCast(*engineInputs.draftProbs); - params.outputPackedMask = bufferCast(*engineInputs.packedMasks); - params.outputPositionOffsets = bufferCast(*engineInputs.positionOffsets); - params.outputPositionIds = bufferCast(*engineInputs.positionIds); - - params.cumSumGenerationLengths = bufferCast(*cumSumGenerationLengths); - - params.checkParams(); - - // Pack tensors from batch slot position to continuous array - tksd::invokePackGenerationLengths(params, stream.get()); - - if (numGenSequences) - { - // Compute inclusive sum - tksd::invokeScanGenerationLengths(bufferCast(*scanTempStorage), scanTempStorageBytes, - bufferCast(*engineInputs.generationLengths), bufferCast(*cumSumGenerationLengths), - numGenSequences, stream.get()); - } - - // Pack tensors from batch slot position to continuous array - tksd::invokePackExplicitDraftTokens(params, stream.get()); - - if (numGenSequences) - { - // Copy draft probs - tksd::invokeCopyProbs(params, stream.get()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void ExplicitDraftTokensBuffers::setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, - ITensor const& requestTypes, ITensor const& seqSlots, ExplicitDraftTokensBuffers::Inputs const& draftBuffers, - ITensor const& contextPositionIds, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - runtime::BufferManager const& manager, runtime::CudaStream const& stream) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const explicitDraftTokensModule = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const seqSlotsPtr = bufferCast(seqSlots); - auto const generationLengthsPtr = bufferCast(*draftBuffers.generationLengthsHost); - SizeType32 totalGenLengths = 0; - for (SizeType32 si = 0; si < numGenSequences; ++si) - { - auto const slot = seqSlotsPtr[numCtxSequences + si]; - totalGenLengths += generationLengthsPtr[slot]; - } - - // Reshape position ids. - engineInputs.positionIds->reshape(ITensor::makeShape({contextPositionIds.getShape().d[0] + totalGenLengths})); - // Copy position ids -- hacky solution to avoid filling them for the context requests. - TensorPtr posIdsSlice = ITensor::slice(engineInputs.positionIds, 0, contextPositionIds.getShape().d[0]); - manager.copy(contextPositionIds, *posIdsSlice); - - manager.copy(requestTypes, *engineInputs.requestTypesDevice); - - auto const numSequences = numCtxSequences + numGenSequences; - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const dtype = modelConfig.getDataType(); - - switch (dtype) - { - case tensorrt_llm::DataType::kFLOAT: - setFromInputs(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, - contextPositionIds, *explicitDraftTokensModule, stream); - break; - case tensorrt_llm::DataType::kHALF: - setFromInputs(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, - contextPositionIds, *explicitDraftTokensModule, stream); - break; - case tensorrt_llm::DataType::kBF16: - setFromInputs<__nv_bfloat16>(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, - contextPositionIds, *explicitDraftTokensModule, stream); - break; - default: - TLLM_THROW("DataType %d not supported in ExplicitDraftTokensBuffers", static_cast(dtype)); - break; - } - - // reshape outputs - auto draftTokensShape = engineOutputs.nextDraftTokens->getShape(); - draftTokensShape.d[0] = numSequences; - engineOutputs.nextDraftTokens->reshape(draftTokensShape); - auto draftIndicesShape = engineOutputs.nextDraftIndices->getShape(); - draftIndicesShape.d[0] = numSequences; - engineOutputs.nextDraftIndices->reshape(draftIndicesShape); - auto draftProbsShape = engineOutputs.nextDraftProbs->getShape(); - draftProbsShape.d[0] = numSequences; - engineOutputs.nextDraftProbs->reshape(draftProbsShape); - - auto maxGenLength = bufferCast(*draftBuffers.maxGenLengthHost)[0]; - if (maxGenLength == 0) - { - maxGenLength = explicitDraftTokensModule->getMaxDecodingTokens(); - } - auto positionOffsetsShape = engineInputs.positionOffsets->getShape(); - positionOffsetsShape.d[1] = maxGenLength; - engineInputs.positionOffsets->reshape(positionOffsetsShape); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void ExplicitDraftTokensBuffers::insertInputTensors( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::WorldConfig const& /* worldConfig */) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - // inputs - inputBuffers.insert_or_assign("redrafter_inverted_temperature", engineInputs.temperatures); - inputBuffers.insert_or_assign("device_request_types", engineInputs.requestTypesDevice); - - inputBuffers.insert_or_assign("spec_decoding_generation_lengths", engineInputs.generationLengths); - inputBuffers.insert_or_assign("spec_decoding_position_offsets", engineInputs.positionOffsets); - inputBuffers.insert_or_assign("spec_decoding_packed_mask", engineInputs.packedMasks); - - inputBuffers.insert_or_assign("draft_tokens", engineInputs.draftTokens); - inputBuffers.insert_or_assign("draft_indices", engineInputs.draftIndices); - inputBuffers.insert_or_assign("draft_probs", engineInputs.draftProbs); - - inputBuffers.insert_or_assign("rand_data_sample", engineInputs.randomDataSample); - inputBuffers.insert_or_assign("rand_data_validation", engineInputs.randomDataValidation); - inputBuffers.insert_or_assign("position_ids_base", engineInputs.positionIdsBase); - inputBuffers.insert_or_assign("position_ids", engineInputs.positionIds); - inputBuffers.insert_or_assign("spec_decoding_use", engineInputs.useSpecDecoding); - - // outputs - outputBuffers.insert_or_assign("next_spec_decoding_generation_lengths", engineOutputs.nextGenerationLengths); - outputBuffers.insert_or_assign("next_spec_decoding_position_offsets", engineOutputs.nextPositionOffsets); - outputBuffers.insert_or_assign("spec_decoding_mask", engineOutputs.masks); - - outputBuffers.insert_or_assign("next_draft_tokens", engineOutputs.nextDraftTokens); - outputBuffers.insert_or_assign("next_draft_indices", engineOutputs.nextDraftIndices); - outputBuffers.insert_or_assign("next_draft_probs", engineOutputs.nextDraftProbs); - outputBuffers.insert_or_assign("next_flat_tokens", engineOutputs.nextFlatTokens); - - outputBuffers.insert_or_assign("num_accepted_tokens", engineOutputs.bestPathLengths); - outputBuffers.insert_or_assign("accepted_beam_index", engineOutputs.bestPathIndices); - outputBuffers.insert_or_assign("max_gen_token", engineOutputs.maxGenToken); - outputBuffers.insert_or_assign("total_gen_token", engineOutputs.totalGenToken); - outputBuffers.insert_or_assign("packed_position_ids", engineOutputs.packedPositionIds); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/gptDecoder.cpp b/cpp/tensorrt_llm/runtime/gptDecoder.cpp deleted file mode 100644 index e1ac1717af45..000000000000 --- a/cpp/tensorrt_llm/runtime/gptDecoder.cpp +++ /dev/null @@ -1,767 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/gptDecoder.h" - -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/layers/dynamicDecodeLayer.h" -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" - -#include "tensorrt_llm/common/tllmDataType.h" - -#include - -namespace tle = tensorrt_llm::executor; -namespace tl = tensorrt_llm::layers; - -using namespace tensorrt_llm::runtime; - -using BufferConstPtr = IBuffer::SharedConstPtr; -using BufferPtr = IBuffer::SharedPtr; -using TensorConstPtr = ITensor::SharedConstPtr; -using TensorPtr = ITensor::SharedPtr; - -template -GptDecoder::GptDecoder(executor::DecodingMode const& mode, size_t maxNumSequences, size_t maxBeamWidth, - size_t vocabSize, size_t vocabSizePadded, CudaStreamPtr const& stream, - std::shared_ptr speculativeDecodingModule) - : mManager{std::make_shared(stream)} - , mMaxNumSequences(maxNumSequences) - , mVocabSize(vocabSize) - , mVocabSizePadded(vocabSizePadded) - , mDecodingMode{mode} -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain( - maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, speculativeDecodingModule); - mDynamicDecodeLayer = std::make_shared>(mode, decodingDomain, mManager); - - mDecodingLayerWorkspace = std::make_unique( - mManager, decodingDomain, TRTDataType::value, mDynamicDecodeLayer->getWorkspaceSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void GptDecoder::disableLookahead( - std::optional const& samplingConfig, SizeType32 batchSize, TensorConstPtr batchSlots) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mDecodingMode = executor::DecodingMode::TopKTopP(); - auto const decodingDomain - = tensorrt_llm::layers::DecoderDomain(mMaxNumSequences, 1, mVocabSize, mVocabSizePadded, nullptr); - - auto setupParams = std::make_shared(); - - if (batchSize == 0) - { - mDynamicDecodeLayer->disableLookahead( - decodingDomain, batchSize, batchSlots, setupParams, mDecodingLayerWorkspace); - return; - } - - mSamplingConfig = samplingConfig.value(); - TLLM_CHECK_WITH_INFO(mSamplingConfig.validate(), "Sampling config is invalid"); - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots are mandatory to set up the decoder."); - // penalty parameters - auto penaltyParams = std::make_shared(); - penaltyParams->repetitionPenalty = mSamplingConfig.repetitionPenalty; - penaltyParams->presencePenalty = mSamplingConfig.presencePenalty; - penaltyParams->frequencyPenalty = mSamplingConfig.frequencyPenalty; - penaltyParams->promptIgnoreLength = mSamplingConfig.promptIgnoreLength; - penaltyParams->temperature = mSamplingConfig.temperature; - penaltyParams->minLength = mSamplingConfig.minLength; - - // banwords parameters - auto banWordsParams = std::make_shared(); - banWordsParams->noRepeatNgramSize = mSamplingConfig.noRepeatNgramSize; - - // sampling parameters - auto samplingParams = std::make_shared(); - samplingParams->normalizeLogProbs = mSamplingConfig.normalizeLogProbs; - if (mSamplingConfig.topK) - { - auto const& topK = mSamplingConfig.topK.value(); - samplingParams->runtimeTopK = std::vector(std::begin(topK), std::end(topK)); - } - samplingParams->runtimeTopP = mSamplingConfig.topP; - samplingParams->topPDecay = mSamplingConfig.topPDecay; - samplingParams->topPMin = mSamplingConfig.topPMin; - samplingParams->topPResetIds = mSamplingConfig.topPResetIds; - samplingParams->outputLogProbs = mSamplingConfig.outputLogProbs; - samplingParams->cumLogProbs = mSamplingConfig.cumLogProbs; - samplingParams->runtimeMinP = mSamplingConfig.minP; - - // get setup parameters - setupParams->penaltyParams = std::move(penaltyParams); - setupParams->banWordsParams = std::move(banWordsParams); - setupParams->decodingParams = std::move(samplingParams); - - mDecodingLayerWorkspace->setDeviceBatchSlots(batchSlots); - mDynamicDecodeLayer->disableLookahead(decodingDomain, batchSize, batchSlots, setupParams, mDecodingLayerWorkspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void GptDecoder::setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, - std::optional const& output, std::optional explicitDraftTokensDType, - std::optional> const& lookaheadPrompt, - std::optional> const& lookaheadAlgoConfigs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mSamplingConfig = samplingConfig; - auto setupParams = std::make_shared(); - - TLLM_CHECK_WITH_INFO(mSamplingConfig.validate(), "Sampling config is invalid"); - TLLM_CHECK_WITH_INFO(batchSlots != nullptr, "Batch slots are mandatory to set up the decoder."); - - auto penaltyParams = std::make_shared(); - penaltyParams->repetitionPenalty = mSamplingConfig.repetitionPenalty; - penaltyParams->presencePenalty = mSamplingConfig.presencePenalty; - penaltyParams->frequencyPenalty = mSamplingConfig.frequencyPenalty; - penaltyParams->promptIgnoreLength = mSamplingConfig.promptIgnoreLength; - penaltyParams->temperature = mSamplingConfig.temperature; - penaltyParams->minLength = mSamplingConfig.minLength; - - setupParams->penaltyParams = std::move(penaltyParams); - - auto banWordsParams = std::make_shared(); - banWordsParams->noRepeatNgramSize = mSamplingConfig.noRepeatNgramSize; - - setupParams->banWordsParams = std::move(banWordsParams); - - if (mDecodingMode.isTopKorTopP()) - { - auto samplingParams = std::make_shared(); - samplingParams->normalizeLogProbs = mSamplingConfig.normalizeLogProbs; - // signed to unsigned - if (mSamplingConfig.topK) - { - auto const& topK = mSamplingConfig.topK.value(); - samplingParams->runtimeTopK = std::vector(std::begin(topK), std::end(topK)); - } - - samplingParams->runtimeTopP = mSamplingConfig.topP; - samplingParams->topPDecay = mSamplingConfig.topPDecay; - samplingParams->topPMin = mSamplingConfig.topPMin; - samplingParams->topPResetIds = mSamplingConfig.topPResetIds; - samplingParams->outputLogProbs = mSamplingConfig.outputLogProbs; - samplingParams->cumLogProbs = mSamplingConfig.cumLogProbs; - samplingParams->runtimeMinP = mSamplingConfig.minP; - - setupParams->decodingParams = std::move(samplingParams); - } - else if (mDecodingMode.isBeamSearch()) - { - auto beamSearchParams = std::make_shared(); - beamSearchParams->beamSearchDiversityRate = mSamplingConfig.beamSearchDiversityRate; - beamSearchParams->lengthPenalty = mSamplingConfig.lengthPenalty; - beamSearchParams->earlyStopping = mSamplingConfig.earlyStopping; - beamSearchParams->beamWidthArray = mSamplingConfig.beamWidthArray; - - setupParams->decodingParams = std::move(beamSearchParams); - } - else if (mDecodingMode.isMedusa()) - { - auto medusaParams = std::make_shared(); - // signed to unsigned - if (mSamplingConfig.topK) - { - auto const& topK = mSamplingConfig.topK.value(); - medusaParams->runtimeTopK = std::vector(std::begin(topK), std::end(topK)); - } - medusaParams->runtimeHeadsTopK = mSamplingConfig.topKMedusaHeads; - - setupParams->decodingParams = std::move(medusaParams); - } - else if (mDecodingMode.isExplicitDraftTokens()) - { - TLLM_CHECK_WITH_INFO(output.has_value(), "Output tensors must be provided for ExplicitDraftTokens"); - auto explicitDraftTokensParams = std::make_shared(); - explicitDraftTokensParams->temperature = mSamplingConfig.temperature; - explicitDraftTokensParams->randomDataSample = output->explicitDraftTokensBuffers->randomDataSample; - explicitDraftTokensParams->temperatures = output->explicitDraftTokensBuffers->temperatures; - TLLM_CHECK(explicitDraftTokensDType.has_value()); - explicitDraftTokensParams->dtype = explicitDraftTokensDType.value(); - - setupParams->decodingParams = explicitDraftTokensParams; - } - else if (mDecodingMode.isLookahead()) - { - TLLM_LOG_DEBUG("gptDecoder setup lookahead, batchSize=%d", batchSize); - auto lookaheadParams = std::make_shared(); - - TLLM_CHECK_WITH_INFO(lookaheadPrompt.has_value(), "Lookahead prompt must be provided"); - lookaheadParams->prompt = lookaheadPrompt.value(); - TLLM_CHECK_WITH_INFO(lookaheadAlgoConfigs.has_value(), "Lookahead algo configs must be provided"); - lookaheadParams->algoConfigs = lookaheadAlgoConfigs.value(); - TLLM_CHECK_WITH_INFO(output.has_value(), "Output tensors must be provided for Lookahead decoding"); - lookaheadParams->generationLengths = output->lookaheadOutputs->generationLengths; - lookaheadParams->positionOffsets = output->lookaheadOutputs->positionOffsets; - lookaheadParams->attentionPackedMasks = output->lookaheadOutputs->packedMasks; - - setupParams->decodingParams = std::move(lookaheadParams); - } - else if (mDecodingMode.isExternalDraftTokens()) - { - auto externalDraftTokensParams = std::make_shared(); - // signed to unsigned - if (mSamplingConfig.topK) - { - auto const& topK = mSamplingConfig.topK.value(); - externalDraftTokensParams->runtimeTopK = std::vector(std::begin(topK), std::end(topK)); - } - externalDraftTokensParams->runtimeTopP = mSamplingConfig.topP; - setupParams->decodingParams = std::move(externalDraftTokensParams); - } - else if (mDecodingMode.isEagle()) - { - TLLM_CHECK_WITH_INFO(output.has_value(), "Output tensors must be provided for Eagle"); - auto eagleParams = std::make_shared(); - eagleParams->temperature = mSamplingConfig.originalTemperature; - eagleParams->randomDataSample = output->eagleBuffers->randomDataSample; - eagleParams->temperatures = output->eagleBuffers->temperatures; - - setupParams->decodingParams = eagleParams; - } - - setupParams->decodingParams->randomSeed = mSamplingConfig.randomSeed; - - mDynamicDecodeLayer->setup(batchSize, mSamplingConfig.beamWidth, batchSlots, setupParams, mDecodingLayerWorkspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ - -std::shared_ptr prepareBanWordsInputs(DecodingInput const& input) -{ - auto banWordsParams = std::make_shared(input.batchSize); - if (input.badWordsPtrs) - { - TLLM_CHECK_WITH_INFO(input.badWordsPtrs, "Bad word lengths must be provided when badWordsPtrs is given"); - banWordsParams->badWordsPtr = input.badWordsPtrs; - banWordsParams->badWordsLengths = input.badWordsLens; - banWordsParams->maxBadWordsLen = input.maxBadWordsLen; - } - - return banWordsParams; -} - -std::shared_ptr prepareStopCriteriaInputs(DecodingInput const& input) -{ - auto stopCriteriaParams = std::make_shared(input.batchSize); - if (input.stopWordsPtrs) - { - TLLM_CHECK_WITH_INFO(input.stopWordsLens, "Stop word lengths must be provided when stopWordsPtrs is given"); - - stopCriteriaParams->stopWordsPtr = input.stopWordsPtrs; - stopCriteriaParams->stopWordsLengths = input.stopWordsLens; - stopCriteriaParams->maxStopWordsLen = input.maxStopWordsLen; - } - - if (input.sequenceLimitLength) - { - stopCriteriaParams->sequenceLimitLength = input.sequenceLimitLength; - } - - return stopCriteriaParams; -} - -void prepareMedusaInputs( - DecodingInput const& inputs, size_t maxNumSequences, std::shared_ptr& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto inputParams = std::dynamic_pointer_cast(baseInputs); - - auto const& medusaInputs = inputs.medusaInputs.value(); - - inputParams->curTokensPerStep = medusaInputs.medusaCurTokensPerStep; - inputParams->targetTokensPerStep = medusaInputs.medusaTargetTokensPerStep; - inputParams->paths = medusaInputs.medusaPaths; - inputParams->treeIds = medusaInputs.medusaTreeIds; - auto const batchSlots = bufferCast(*inputs.batchSlots); - if (medusaInputs.medusaLogits.size()) - { - std::vector> medusaLogits; - auto const batchSize = medusaInputs.medusaLogits.size(); - medusaLogits.resize(maxNumSequences); - for (size_t bi = 0; bi < batchSize; ++bi) - { - auto const slot = batchSlots[bi]; - auto const& logitsHeads = medusaInputs.medusaLogits.at(slot); - auto const medusaHeads = logitsHeads.size(); - medusaLogits[slot].resize(medusaHeads); - for (size_t hi = 0; hi < medusaHeads; ++hi) - { - if (logitsHeads[hi]) - { - medusaLogits[slot][hi] = logitsHeads[hi]; - } - } - } - inputParams->medusaLogits = medusaLogits; - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void prepareExternalDraftTokensInputs(DecodingInput const& inputs, std::shared_ptr& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto inputParams = std::dynamic_pointer_cast(baseInputs); - auto const& externalDraftTokensInputs = inputs.externalDraftTokensInputs.value(); - - inputParams->draftLogits = externalDraftTokensInputs.draftLogits; - inputParams->draftProbs = externalDraftTokensInputs.draftProbs; - inputParams->targetProbs = externalDraftTokensInputs.targetProbs; - inputParams->numDraftTokens = externalDraftTokensInputs.numDraftTokens; - inputParams->numDraftTokensHost = externalDraftTokensInputs.numDraftTokensHost; - inputParams->draftTokenIds = externalDraftTokensInputs.draftTokenIds; - inputParams->constantThreshold = externalDraftTokensInputs.constantThreshold; - inputParams->useRandomAcceptanceThreshold = externalDraftTokensInputs.useRandomAcceptanceThreshold; - inputParams->step = externalDraftTokensInputs.step; - inputParams->useDraftLogits = externalDraftTokensInputs.useDraftLogits; - inputParams->useDraftLogitsHost = externalDraftTokensInputs.useDraftLogitsHost; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void prepareExplicitDraftTokensInput(DecodingInput const& inputs, std::shared_ptr& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto inputParams = std::dynamic_pointer_cast(baseInputs); - - auto& explicitDraftTokensInputs = inputs.explicitDraftTokensInputs; - - TLLM_CHECK_WITH_INFO(explicitDraftTokensInputs.has_value(), "ExplicitDraftTokensInputs are not set"); - - inputParams->nextDraftTokens = explicitDraftTokensInputs->nextDraftTokens; - inputParams->nextFlatTokens = explicitDraftTokensInputs->nextFlatTokens; - inputParams->nextDraftIndices = explicitDraftTokensInputs->nextDraftIndices; - inputParams->nextDraftProbs = explicitDraftTokensInputs->nextDraftProbs; - inputParams->lastDraftTokens = explicitDraftTokensInputs->lastDraftTokens; - inputParams->lastDraftIndices = explicitDraftTokensInputs->lastDraftIndices; - inputParams->masks = explicitDraftTokensInputs->masks; - inputParams->packedPosIds = explicitDraftTokensInputs->packedPositionIds; - inputParams->bestPathLengths = explicitDraftTokensInputs->bestPathLengths; - inputParams->bestPathIndices = explicitDraftTokensInputs->bestPathIndices; - inputParams->generationLengths = explicitDraftTokensInputs->nextGenerationLengths; - inputParams->positionIdsBase = explicitDraftTokensInputs->lastPositionIdsBase; - inputParams->lastGenerationLengths = explicitDraftTokensInputs->lastGenerationLengths; - inputParams->maxGenLengthDevice = explicitDraftTokensInputs->maxGenLengthDevice; - inputParams->seqSlots = explicitDraftTokensInputs->seqSlots; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void prepareLookaheadInputs(DecodingInput const& inputs, std::shared_ptr& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto inputParams = std::dynamic_pointer_cast(baseInputs); - auto const& lookaheadInputs = inputs.lookaheadInputs.value(); - inputParams->curTokensPerStep = lookaheadInputs.tokensPerStep; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void prepareEagleInput(DecodingInput const& inputs, std::shared_ptr& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto inputParams = std::dynamic_pointer_cast(baseInputs); - - auto& eagleInputs = inputs.eagleInputs; - - TLLM_CHECK_WITH_INFO(eagleInputs.has_value(), "EagleInputs are not set"); - - inputParams->nextDraftTokens = eagleInputs->nextDraftTokens; - inputParams->nextDraftLens = eagleInputs->nextDraftLens; - inputParams->nextDraftPaths = eagleInputs->nextDraftPaths; - inputParams->lastDraftTokens = eagleInputs->lastDraftTokens; - inputParams->lastDraftLens = eagleInputs->lastDraftLens; - inputParams->lastDraftPaths = eagleInputs->lastDraftPaths; - inputParams->acceptedTokens = eagleInputs->acceptedTokens; - inputParams->acceptedLens = eagleInputs->acceptedLens; - inputParams->acceptedPathIds = eagleInputs->acceptedPathIds; - inputParams->chunkedContextNextTokens = eagleInputs->chunkedContextNextTokens; - inputParams->seqSlots = eagleInputs->seqSlots; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -std::shared_ptr prepareInputs( - DecodingInput const& input, size_t maxNumSequences, tle::DecodingMode const& decodingMode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto constexpr ite = 0; - - TLLM_CHECK_WITH_INFO(input.batchSlots != nullptr, "Batch slots are mandatory to call the decoder."); - std::shared_ptr forwardParams; - if (decodingMode.isTopKorTopP()) - { - forwardParams - = std::make_shared(input.endIds, input.batchSlots, input.step, ite, input.batchSize); - } - else if (decodingMode.isBeamSearch()) - { - forwardParams = std::make_shared(input.endIds, input.batchSlots, input.step, ite, - input.batchSize, input.maxAttentionWindow, input.sinkTokenLength); - - if (input.cacheIndirection) - { - forwardParams->srcCacheIndirection = input.cacheIndirection; - } - forwardParams->beamSearchSteps = input.generationSteps; - } - else if (decodingMode.isMedusa()) - { - forwardParams = std::make_shared(input.endIds, input.batchSlots, input.batchSize); - } - else if (decodingMode.isLookahead()) - { - forwardParams = std::make_shared(input.endIds, input.batchSlots); - } - else if (decodingMode.isExplicitDraftTokens()) - { - forwardParams - = std::make_shared(input.endIds, input.batchSlots, input.batchSize); - } - else if (decodingMode.isExternalDraftTokens()) - { - forwardParams = std::make_shared( - input.endIds, input.batchSlots, input.step, ite, input.batchSize); - } - else if (decodingMode.isEagle()) - { - auto& eagleInputs = input.eagleInputs; - - TLLM_CHECK_WITH_INFO(eagleInputs.has_value(), "EagleInputs are not set"); - - forwardParams = std::make_shared(input.endIds, input.batchSlots, input.batchSize, - eagleInputs->nextDraftTokens, eagleInputs->nextDraftLens, eagleInputs->nextDraftPaths, - eagleInputs->lastDraftTokens, eagleInputs->lastDraftLens, eagleInputs->lastDraftPaths, - eagleInputs->acceptedTokens, eagleInputs->acceptedLens, eagleInputs->acceptedPathIds, - eagleInputs->chunkedContextNextTokens, eagleInputs->seqSlots); - } - - // No logits for explicit draft tokens and eagle - if (!decodingMode.isExplicitDraftTokens() && !decodingMode.isEagle()) - { - for (auto const& logits : input.logitsVec) - { - TLLM_CHECK(logits->getDataType() == TRTDataType::value); - } - forwardParams->logitsVec = input.logitsVec; - } - - if (input.embeddingBias) - { - forwardParams->embeddingBias = input.embeddingBias; - } - - if (input.lengths) - { - forwardParams->inputLengths = input.lengths; - } - - forwardParams->banWordsInputs = prepareBanWordsInputs(input); - - forwardParams->stopCriteriaInputs = prepareStopCriteriaInputs(input); - - if (input.finishReasons) - { - forwardParams->finished = input.finishReasons; - } - - // Speculative decoding - if (decodingMode.isMedusa()) - { - prepareMedusaInputs(input, maxNumSequences, forwardParams); - } - else if (decodingMode.isExplicitDraftTokens()) - { - prepareExplicitDraftTokensInput(input, forwardParams); - } - else if (decodingMode.isLookahead() && input.lookaheadInputs) - { - prepareLookaheadInputs(input, forwardParams); - forwardParams->localBatchSize = input.batchSize; - } - else if (decodingMode.isExternalDraftTokens()) - { - prepareExternalDraftTokensInputs(input, forwardParams); - } - else if (decodingMode.isEagle()) - { - prepareEagleInput(input, forwardParams); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return forwardParams; -} - -void prepareBeamSearchOutputs(DecodingOutput& output, std::shared_ptr& baseOutputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const& bhSrc = output.beamHypotheses; - auto bhOutputs = std::dynamic_pointer_cast(baseOutputs); - bhOutputs->beamHypotheses = std::make_unique(); - auto& bhDst = bhOutputs->beamHypotheses; - - if (bhSrc.outputIdsCBA) - { - bhDst->outputIdsCBA = bufferCast(*bhSrc.outputIdsCBA); - } - if (bhSrc.logProbsCBA) - { - bhDst->logProbsCBA = bufferCast(*bhSrc.logProbsCBA); - } - if (bhSrc.sequenceLengthsCBA) - { - bhDst->sequenceLengthsCBA = bufferCast(*bhSrc.sequenceLengthsCBA); - } - if (bhSrc.cumLogProbsCBA) - { - bhDst->cumLogProbsCBA = bufferCast(*bhSrc.cumLogProbsCBA); - } - if (bhSrc.normedScoresCBA) - { - bhDst->normedScoresCBA = bufferCast(*bhSrc.normedScoresCBA); - } - if (bhSrc.numBeamsCBA) - { - bhDst->numBeamsCBA = bufferCast(*bhSrc.numBeamsCBA); - } - if (bhSrc.minNormedScoresCBA) - { - bhDst->minNormedScoresCBA = bufferCast(*bhSrc.minNormedScoresCBA); - } - if (bhSrc.batchDones) - { - bhDst->batchDones = bufferCast(*bhSrc.batchDones); - } - if (output.cacheIndirection) - { - bhOutputs->tgtCacheIndirection = output.cacheIndirection; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void prepareSpeculativeDecodingOutputs(DecodingOutput& output, std::shared_ptr& baseOutputs, - tle::DecodingMode const& decodingMode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto outputParams = std::dynamic_pointer_cast(baseOutputs); - - auto const& speculativeDecodingOutputs = output.speculativeDecodingOutputs; - TLLM_CHECK_WITH_INFO(speculativeDecodingOutputs.has_value(), "speculativeDecodingOutputs is not set"); - - outputParams->nextDraftTokens = speculativeDecodingOutputs->nextDraftTokens; - outputParams->numNewTokens = speculativeDecodingOutputs->acceptedTokensLen; - outputParams->numNewTokensCumSum = speculativeDecodingOutputs->acceptedLengthsCumSum; - outputParams->pathsOffsets = speculativeDecodingOutputs->pathsOffsets; - if (speculativeDecodingOutputs->nextDraftTokensLen) - { - outputParams->nextDraftLengths = speculativeDecodingOutputs->nextDraftTokensLen; - } - if (speculativeDecodingOutputs->prevDraftTokensLen) - { - outputParams->prevDraftLengths = speculativeDecodingOutputs->prevDraftTokensLen; - } - - if (decodingMode.isExplicitDraftTokens()) - { - auto outputParams = std::dynamic_pointer_cast(baseOutputs); - auto const& explicitDraftTokensBuffers = output.explicitDraftTokensBuffers; - TLLM_CHECK_WITH_INFO(explicitDraftTokensBuffers.has_value(), "explicitDraftTokensBuffers is not set"); - outputParams->packedMasks = explicitDraftTokensBuffers->packedMasks; - outputParams->nextDraftPosIds = explicitDraftTokensBuffers->positionIds; - - outputParams->unpackedNextDraftTokens = explicitDraftTokensBuffers->draftTokens; - outputParams->unpackedNextDraftIndices = explicitDraftTokensBuffers->draftIndices; - outputParams->nextDraftProbs = explicitDraftTokensBuffers->draftProbs; - outputParams->positionIdsBase = explicitDraftTokensBuffers->positionIdsBase; - outputParams->randomDataSample = explicitDraftTokensBuffers->randomDataSample; - outputParams->randomDataValidation = explicitDraftTokensBuffers->randomDataValidation; - outputParams->temperatures = explicitDraftTokensBuffers->temperatures; - outputParams->generationLengths = explicitDraftTokensBuffers->generationLengths; - outputParams->generationLengthsHost = explicitDraftTokensBuffers->generationLengthsHost; - outputParams->maxGenLengthHost = explicitDraftTokensBuffers->maxGenLengthHost; - } - else if (decodingMode.isLookahead()) - { - TLLM_CHECK(output.lookaheadOutputs); - auto outputParams = std::dynamic_pointer_cast(baseOutputs); - outputParams->packedMasks = output.lookaheadOutputs->packedMasks; - outputParams->positionIds = output.lookaheadOutputs->positionIds; - outputParams->positionOffsets = output.lookaheadOutputs->positionOffsets; - outputParams->generationLengths = output.lookaheadOutputs->generationLengths; - } - else if (decodingMode.isEagle()) - { - auto outputParams = std::dynamic_pointer_cast(baseOutputs); - auto const& eagleBuffers = output.eagleBuffers; - TLLM_CHECK_WITH_INFO(eagleBuffers.has_value(), "eagleBuffers is not set"); - - outputParams->temperatures = eagleBuffers->temperatures; - outputParams->unpackedNextDraftTokens = eagleBuffers->draftTokens; - outputParams->nextDraftPaths = eagleBuffers->draftPaths; - outputParams->generationLengths = eagleBuffers->specDecodingGenerationLengths; - outputParams->generationLengthsHost = eagleBuffers->specDecodingGenerationLengthsHost; - outputParams->nextDraftPosIds = eagleBuffers->specDecodingPositionOffsets; - outputParams->packedMasks = eagleBuffers->specDecodingPackedMasks; - outputParams->randomDataSample = eagleBuffers->randomDataSample; - outputParams->randomDataValidation = eagleBuffers->randomDataValidation; - - outputParams->eagleNetCtxRequestTypesHost = eagleBuffers->eagleNetCtxRequestTypesHost; - outputParams->eagleNetCtxContextLengthsHost = eagleBuffers->eagleNetCtxContextLengthsHost; - outputParams->eagleNetCtxPastKeyValueLengthsHost = eagleBuffers->eagleNetCtxPastKeyValueLengthsHost; - outputParams->eagleNetGenRequestTypesHost = eagleBuffers->eagleNetGenRequestTypesHost; - outputParams->eagleNetGenContextLengthsHost = eagleBuffers->eagleNetGenContextLengthsHost; - outputParams->eagleNetGenPastKeyValueLengthsHost = eagleBuffers->eagleNetGenPastKeyValueLengthsHost; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::shared_ptr prepareOutputs(DecodingOutput& output, tle::DecodingMode const& decodingMode) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - std::shared_ptr outputParams; - - if (decodingMode.isBeamSearch()) - { - outputParams = std::make_shared(output.ids); - } - else if (decodingMode.isMedusa()) - { - outputParams = std::make_shared(output.ids); - } - else if (decodingMode.isLookahead()) - { - outputParams = std::make_shared(output.ids); - } - else if (decodingMode.isExplicitDraftTokens()) - { - outputParams = std::make_shared(output.ids); - } - else if (decodingMode.isEagle()) - { - outputParams = std::make_shared(output.ids); - } - else - { - outputParams = std::make_shared(output.ids); - } - - // Common outputs - outputParams->newTokens = output.newTokens; - - if (output.cumLogProbs) - { - outputParams->cumLogProbs = output.cumLogProbs; - } - - if (output.parentIds) - { - outputParams->parentIds = output.parentIds; - } - - if (output.finishReasons) - { - outputParams->finished = output.finishReasons; - } - - if (output.finishedSum) - { - outputParams->finishedSum = output.finishedSum; - } - - if (output.lengths) - { - outputParams->sequenceLength = output.lengths; - } - - if (output.logProbs) - { - outputParams->outputLogProbs = output.logProbs; - outputParams->outputLogProbsTiled = output.logProbsTiled; - } - - // Beam search outputs - if (decodingMode.isBeamSearch()) - { - prepareBeamSearchOutputs(output, outputParams); - } - - // Speculative decoding outputs - if (decodingMode.isMedusa() || decodingMode.isLookahead() || decodingMode.isExplicitDraftTokens() - || decodingMode.isEagle()) - { - prepareSpeculativeDecodingOutputs(output, outputParams, decodingMode); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return outputParams; -} - -} // namespace - -template -void GptDecoder::forwardAsync(DecodingOutput& output, DecodingInput const& input) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto forwardParams = prepareInputs(input, mMaxNumSequences, mDecodingMode); - auto outputParams = prepareOutputs(output, mDecodingMode); - mDynamicDecodeLayer->forwardAsync(outputParams, forwardParams, mDecodingLayerWorkspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void GptDecoder::forwardSync(DecodingOutput& output, DecodingInput const& input) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto forwardParams = prepareInputs(input, mMaxNumSequences, mDecodingMode); - auto outputParams = prepareOutputs(output, mDecodingMode); - - mDynamicDecodeLayer->forwardSync(outputParams, forwardParams, mDecodingLayerWorkspace); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace tensorrt_llm::runtime -{ -template class GptDecoder; -template class GptDecoder; -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp deleted file mode 100644 index 7b3a12ed7a2c..000000000000 --- a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp +++ /dev/null @@ -1,257 +0,0 @@ -/* - * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/gptDecoderBatched.h" - -#include "common.h" -#include "decoderState.h" -#include "iBuffer.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/kernels/decodingKernels.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaEvent.h" - -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; -namespace tb = tensorrt_llm::batch_manager; -using TensorPtr = ITensor::SharedPtr; - -GptDecoderBatched::GptDecoderBatched(GptDecoderBatched::CudaStreamPtr stream) - : mRuntimeStream{std::move(stream)} - , mBufferManager{mRuntimeStream} -{ -} - -void GptDecoderBatched::disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - std::vector samplingConfigs; - samplingConfigs.reserve(genRequests.size()); - auto batchSlotsRange = BufferRange(*batchSlots); - - SizeType32 batchIdx = 0; - for (auto const& llmReq : genRequests) - { - samplingConfigs.push_back(llmReq->mSamplingConfig); - batchSlotsRange[batchIdx] = llmReq->mSeqSlot.value(); - batchIdx += 1; - } - auto const batchSize = batchIdx; - std::optional samplingConfig; - if (batchSize > 0) - { - samplingConfig = SamplingConfig(samplingConfigs); - } - TensorPtr batchSlotsView = ITensor::slice(batchSlots, 0, batchSize); - mDecoder->disableLookahead(samplingConfig, batchSize, batchSlots); - - CudaEvent event{}; - mDecoderStream->record(event); - mRuntimeStream->wait(event); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void GptDecoderBatched::setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(maxNumSequences > 0); - TLLM_CHECK(maxBeamWidth > 0); - - std::shared_ptr speculativeDecodingModulePtr = nullptr; - if (modelConfig.getSpeculativeDecodingMode().predictsDraftTokens()) - { - speculativeDecodingModulePtr = modelConfig.getSpeculativeDecodingModulePtr(); - } - - auto const device = mRuntimeStream->getDevice(); - mDecoderStream = std::make_shared(); - TLLM_CHECK(mDecoderStream->getDevice() == device); - - auto const vocabSize = modelConfig.getVocabSize(); - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - mDecoder = IGptDecoder::create(mode, dtype, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, - mDecoderStream, speculativeDecodingModulePtr); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ -//! @brief Prepare Input and Output for decoder step. -// TODO: produce new input and output objects -void prepareForward(decoder::DecoderState const& decoderState, SizeType32 step, tb::DecoderInputBuffers const& input, - BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const speculativeDecodingMode = decoderState.getSpeculativeDecodingMode(); - - auto& dInput = decoderState.getJointDecodingInput(); - auto& dOutput = decoderState.getJointDecodingOutput(); - - dInput.batchSlots = input.forwardBatchSlots.at(step); - dInput.batchSize = static_cast(dInput.batchSlots->getSize()); - dInput.logitsVec = input.batchLogits.at(step); - - if (speculativeDecodingMode.isDraftTokensExternal()) - { - dInput.externalDraftTokensInputs->step = step; - - // WAR: reset finished state for generation requests - if (step == 0) - { - auto batchSlotsRange = BufferRange(*dInput.batchSlots); - for (auto batchSlot : batchSlotsRange) - { - ::TensorPtr finishedStepsSlice = ITensor::slice(decoderState.getFinishReasons(), batchSlot, 1); - bufferManager.setZero(*finishedStepsSlice); - } - } - } - - dOutput.newTokens = ITensor::slice(dOutput.newTokensSteps, step, decoderState.getMaxDecodingDecoderTokens()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace - -void GptDecoderBatched::forwardDispatch(decoder::DecoderState const& decoderState, tb::DecoderInputBuffers const& input) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - for (SizeType32 step = 0; step < input.maxDecoderSteps; ++step) - { - BufferManager manager{mDecoderStream}; - prepareForward(decoderState, step, input, manager); - - if (decoderState.getJointDecodingInput().batchSize > 0) - { - mDecoder->forwardAsync(decoderState.getJointDecodingOutput(), decoderState.getJointDecodingInput()); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -CudaEvent GptDecoderBatched::forwardAsync( - decoder::DecoderState const& decoderState, tb::DecoderInputBuffers const& input) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto eventStart = CudaEvent{}; - mRuntimeStream->record(eventStart); - mDecoderStream->wait(eventStart.get()); - - forwardDispatch(decoderState, input); - - CudaEvent event{}; - mDecoderStream->record(event); - mRuntimeStream->wait(event); - - CudaEvent eventStop{}; - mRuntimeStream->record(eventStop); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return eventStop; -} - -void GptDecoderBatched::forward(decoder::DecoderState const& decoderState, tb::DecoderInputBuffers const& input) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto decoderFinishEvent = forwardAsync(decoderState, input); - decoderFinishEvent.synchronize(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ -std::pair prepareGatherTree( - decoder::DecoderState const& decoderState, SizeType32 batchSlot, bool streaming, CudaStream const& stream) -{ - auto& dJointInput = decoderState.getJointDecodingInput(); - auto& dJointOutput = decoderState.getJointDecodingOutput(); - - auto slice = [batchSlot](auto& a, auto const& b) - { - if (b && b->getShape().d[0] > 0) - { - a = ITensor::slice(b, batchSlot, 1); - } - }; - - // Prepare a slice of dJointInput and dJointOutput for gatherTree - DecodingInput dInput{dJointInput}; - slice(dInput.endIds, dJointInput.endIds); - slice(dInput.lengths, dJointInput.lengths); - - DecodingOutput dOutput{ - ITensor::slice(dJointOutput.ids, batchSlot, 1), ITensor::slice(dJointOutput.gatheredIds, batchSlot, 1)}; - dOutput.beamHypotheses = dJointOutput.beamHypotheses.slice(batchSlot, 1); - slice(dOutput.parentIds, dJointOutput.parentIds); - slice(dOutput.cumLogProbs, dJointOutput.cumLogProbs); - slice(dOutput.cacheIndirection, dJointOutput.cacheIndirection); - slice(dOutput.lengths, dJointOutput.lengths); - slice(dOutput.finishReasons, dJointOutput.finishReasons); - slice(dOutput.logProbs, dJointOutput.logProbs); - - dOutput.newTokens = ITensor::view(dJointOutput.newTokens); - TLLM_CHECK(dOutput.newTokens->getShape().d[0] == 1); - dOutput.newTokens->squeeze(0); - dOutput.newTokens = ITensor::slice(dOutput.newTokens, batchSlot, 1); - dOutput.logProbsTiled = dJointOutput.logProbsTiled; - if (streaming) - { - // in case of streaming we shouldn't overwrite the data in beamHypotheses, since the beam search kernels expect - // ungathered data but the kernels in gatherTree write in-place. - // Thus, we need to make a copy of the beamHypotheses - auto const& beamSearchBuffers = decoderState.getBeamSearchBuffers(); - tensorrt_llm::kernels::invokeCopyBeamHypotheses(dOutput.beamHypotheses, beamSearchBuffers.mOutputBeamHypotheses, - *dOutput.cumLogProbs, *beamSearchBuffers.mCumLogProbsTmp, stream, beamSearchBuffers.mNumSMs); - dOutput.beamHypotheses = beamSearchBuffers.mOutputBeamHypotheses; - dOutput.cumLogProbs = beamSearchBuffers.mCumLogProbsTmp; - } - - return {(std::move(dInput)), (std::move(dOutput))}; -} -} // namespace - -// TODO call this at the end of forward if mFinished[i] changes from false to true? -CudaEvent GptDecoderBatched::finalize(decoder::DecoderState const& decoderState, SizeType32 batchSlot, - SamplingConfig const& samplingConfig, bool streaming) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto [dInput, dOutput] = prepareGatherTree(decoderState, batchSlot, streaming, *mRuntimeStream); - - kernels::gatherTree(dOutput, dInput, samplingConfig, *mRuntimeStream, batchSlot); - - CudaEvent event{}; - mRuntimeStream->record(event); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return event; -} diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index cdf0a0228fac..dbfdec859b69 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -57,7 +57,6 @@ add_library( marlinNvfp4MoeMM.cpp marlinRepack.cpp cudaScaledMM.cpp - dynamicDecodeOp.cpp fmhaPackMaskOp.cpp fp8Op.cpp fp8PerTensorScalingTrtllmGenGemm.cpp @@ -90,7 +89,6 @@ add_library( fusedGatedRMSNormQuant.cpp rmsNormFp4Quant.cpp fusedTopkSoftmax.cpp - gatherTreeOp.cpp groupRmsNormOp.cpp helixPostProcessOp.cpp llama4MinLatency.cpp diff --git a/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp b/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp deleted file mode 100644 index 228b2c614ab7..000000000000 --- a/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp +++ /dev/null @@ -1,464 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/thop/dynamicDecodeOp.h" - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/gptDecoder.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/thop/thUtils.h" -#include -#include - -namespace th = torch; - -namespace tle = tensorrt_llm::executor; -namespace tr = tensorrt_llm::runtime; -namespace tl = tensorrt_llm::layers; -namespace tk = tensorrt_llm::kernels; - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -template -FtDynamicDecode::FtDynamicDecode(size_t const maxBatchSize, size_t const maxBeamWidth, size_t const vocabSize, - size_t const vocabSizePadded, int const tensorParaSize, int const pipelineParaSize) -{ - TLLM_CHECK_WITH_INFO(vocabSizePadded % tensorParaSize == 0, - tensorrt_llm::common::fmtstr( - "vocabSize (%ld) is not multiple of tensorParaSize (%d).", vocabSizePadded, tensorParaSize)); - - auto const decodingDomain = tl::DecoderDomain(maxBatchSize, maxBeamWidth, vocabSize, vocabSizePadded); - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - auto const currentDeviceId = c10::cuda::current_device(); - auto cudaStreamPtr = std::make_shared(stream, currentDeviceId, false); - auto bufferManager = std::make_shared(cudaStreamPtr); - - mFinishedSum = bufferManager->pinnedPool( - tr::ITensor::makeShape({static_cast(maxBatchSize)}), tensorrt_llm::DataType::kINT32); - mDynamicDecodeLayer - = std::make_shared>(tle::DecodingMode::Auto(), decodingDomain, bufferManager); - mBatchSlots = tr::getDefaultBatchSlots(maxBatchSize); - mDecodingWorkspace = std::make_unique(bufferManager, decodingDomain, - tensorrt_llm::runtime::TRTDataType::value, mDynamicDecodeLayer->getWorkspaceSize()); -} - -namespace -{ - -template -void safeInsert(th::optional& tensor, std::optional>& arg) -{ - if (tensor.has_value()) - { - auto shape = convert_shape(tensor.value()); - size_t const size = tensorrt_llm::runtime::ITensor::volume(shape); - auto ptr = get_ptr(tensor.value()); - arg = std::vector(ptr, ptr + size); - } -} - -template -void safeUpdate(th::optional& tensor, std::optional& arg) -{ - if (tensor.has_value()) - { - arg = convert_tensor(tensor.value()); - } -} - -template -void safeUpdate(th::optional& tensor, std::optional& arg) -{ - if (tensor.has_value()) - { - arg = convert_tensor(tensor.value()); - } -} - -template -void safeUpdateScalar(th::optional& tensor, std::optional& arg, std::string const& name) -{ - if (tensor.has_value()) - { - auto accessor = tensor->accessor(); - TLLM_CHECK_WITH_INFO(accessor.size(0) == 1, name + " must be a scalar"); - arg = accessor[0]; - } -} - -template -void safeUpdatePtr(th::optional& tensor, T*& ptr) -{ - if (tensor.has_value()) - { - ptr = get_ptr(tensor.value()); - } -} - -} // namespace - -template -void FtDynamicDecode::setup(size_t const batch_size, size_t const beam_width, - th::optional runtime_top_k_opt, th::optional runtime_top_p_opt, - th::optional temperature_opt, th::optional repetition_penalty_opt, - th::optional presence_penalty_opt, th::optional frequency_penalty_opt, - th::optional prompt_ignore_length_opt, th::optional min_length_opt, - th::optional length_penalty_opt, th::optional early_stopping_opt, - th::optional beam_search_diversity_rate_opt, th::optional random_seed_opt, - th::optional top_p_decay_opt, th::optional top_p_min_opt, - th::optional top_p_reset_ids_opt, th::optional no_repeat_ngram_size_opt, - th::optional min_p_opt, bool output_log_probs, bool cum_log_probs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mBeamWidth = beam_width; - - auto setupParams = std::make_shared(); - auto penaltyParams = std::make_shared(); - auto banWordsParams = std::make_shared(); - safeInsert(temperature_opt, penaltyParams->temperature); - safeInsert(repetition_penalty_opt, penaltyParams->repetitionPenalty); - safeInsert(presence_penalty_opt, penaltyParams->presencePenalty); - safeInsert(frequency_penalty_opt, penaltyParams->frequencyPenalty); - safeInsert(prompt_ignore_length_opt, penaltyParams->promptIgnoreLength); - safeInsert(min_length_opt, penaltyParams->minLength); - safeInsert(no_repeat_ngram_size_opt, banWordsParams->noRepeatNgramSize); - if (beam_width == 1) - { - auto decodingParams = std::make_shared(); - safeInsert(runtime_top_k_opt, decodingParams->runtimeTopK); - safeInsert(runtime_top_p_opt, decodingParams->runtimeTopP); - safeInsert(top_p_decay_opt, decodingParams->topPDecay); - safeInsert(top_p_min_opt, decodingParams->topPMin); - safeInsert(top_p_reset_ids_opt, decodingParams->topPResetIds); - safeInsert(min_p_opt, decodingParams->runtimeMinP); - decodingParams->outputLogProbs = std::vector({output_log_probs}); - decodingParams->cumLogProbs = std::vector({cum_log_probs}); - safeInsert(random_seed_opt, decodingParams->randomSeed); - - setupParams->decodingParams = decodingParams; - } - else - { - auto decodingParams = std::make_shared(); - safeInsert(beam_search_diversity_rate_opt, decodingParams->beamSearchDiversityRate); - safeInsert(length_penalty_opt, decodingParams->lengthPenalty); - safeInsert(early_stopping_opt, decodingParams->earlyStopping); - decodingParams->outputLogProbs = std::vector({output_log_probs}); - decodingParams->cumLogProbs = std::vector({cum_log_probs}); - safeInsert(random_seed_opt, decodingParams->randomSeed); - - setupParams->decodingParams = decodingParams; - } - - // TODO: insert "normalizeLogProbs" and "topKMedusaHeads" - - setupParams->penaltyParams = penaltyParams; - setupParams->banWordsParams = banWordsParams; - mDynamicDecodeLayer->setup( - batch_size, beam_width, tr::ITensor::slice(mBatchSlots, 0, batch_size), setupParams, mDecodingWorkspace); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void FtDynamicDecode::forward(th::Tensor const& logits, int const step, int const maxInputLength, - int const maxAttentionWindow, int const sinkTokenLength, uint64_t const ite, int const localBatchSize, - th::Tensor endId, th::optional embeddingBiasOpt, th::optional inputLengthsOpt, - th::optional sequenceLimitLengthOpt, th::optional stopWordsListPtrsOpt, - th::optional stopWordsLensOpt, int32_t const maxStopWordsLen, - th::optional badWordsListPtrsOpt, th::optional badWordsLensOpt, - int32_t const maxBadWordsLen, th::optional srcCacheIndirectionOpt, th::Tensor& outputTokenIds, - th::Tensor& newTokens, th::Tensor& shouldStop, th::optional finishedInput, - th::optional finishedOutput, th::optional sequenceLengthsOpt, - th::optional cumLogProbsOpt, th::optional outputLogProbsOpt, - th::optional outputLogProbsTiledOpt, th::optional parentIdsOpt, - th::optional tgtCacheIndirectionOpt, th::optional beamHypsOutputIdsCbaOpt, - th::optional beamHypsSeqLenCbaOpt, th::optional beamHypsCumLogProbsCbaOpt, - th::optional beamHypsNormedScoresCbaOpt, th::optional beamHypsLogProbsCbaOpt, - th::optional beamHypsMinNormedScoresOpt, th::optional beamHypsNumBeamsOpt, - th::optional beamHypsIsDoneOpt, bool const useBeamHyps) -{ - TLLM_CHECK_WITH_INFO(mBeamWidth.has_value(), "Beam width is not set. setup() must be called before forward()"); - auto const isBeamSearch = mBeamWidth.value() > 1; - - std::shared_ptr forwardParams; - tr::ITensor::SharedConstPtr batchSlotsSlice = tr::ITensor::slice(mBatchSlots, 0, localBatchSize); - if (isBeamSearch) - { - forwardParams = std::make_shared(convert_tensor(endId), batchSlotsSlice, step, - static_cast(ite), localBatchSize, maxAttentionWindow, sinkTokenLength); - } - else - { - forwardParams = std::make_shared( - convert_tensor(endId), batchSlotsSlice, step, static_cast(ite), localBatchSize); - } - - forwardParams->logits = convert_tensor(logits); - forwardParams->stopCriteriaInputs = std::make_shared(localBatchSize); - forwardParams->banWordsInputs = std::make_shared(localBatchSize); - - safeUpdate(embeddingBiasOpt, forwardParams->embeddingBias); - safeUpdate(inputLengthsOpt, forwardParams->inputLengths); - safeUpdate(sequenceLimitLengthOpt, forwardParams->stopCriteriaInputs->sequenceLimitLength); - safeUpdate(stopWordsListPtrsOpt, forwardParams->stopCriteriaInputs->stopWordsPtr); - safeUpdate(stopWordsLensOpt, forwardParams->stopCriteriaInputs->stopWordsLengths); - forwardParams->stopCriteriaInputs->maxStopWordsLen = maxStopWordsLen; - safeUpdate(badWordsListPtrsOpt, forwardParams->banWordsInputs->badWordsPtr); - safeUpdate(badWordsLensOpt, forwardParams->banWordsInputs->badWordsLengths); - forwardParams->banWordsInputs->maxBadWordsLen = maxBadWordsLen; - safeUpdate(srcCacheIndirectionOpt, forwardParams->srcCacheIndirection); - - tr::ITensor::SharedPtr outputIdsConverted = convert_tensor(outputTokenIds); - - std::shared_ptr outputParams; - if (isBeamSearch) - { - outputParams = std::make_shared(outputIdsConverted); - } - else - { - outputParams = std::make_shared(outputIdsConverted); - } - outputParams->newTokens = convert_tensor(newTokens); - safeUpdate(finishedInput, forwardParams->finished); - safeUpdate(finishedOutput, outputParams->finished); - safeUpdate(sequenceLengthsOpt, outputParams->sequenceLength); - safeUpdate(cumLogProbsOpt, outputParams->cumLogProbs); - safeUpdate(outputLogProbsOpt, outputParams->outputLogProbs); - safeUpdate(outputLogProbsTiledOpt, outputParams->outputLogProbsTiled); - safeUpdate(parentIdsOpt, outputParams->parentIds); - - tr::SizeType32* finishedSumHost = nullptr; - if (forwardParams->stopCriteriaInputs->sequenceLimitLength && outputParams->finished.has_value()) - { - // Skip the initialization and later calculation if there is no limit of sequence length or no finished beam - outputParams->finishedSum = mFinishedSum; - finishedSumHost = tr::bufferCast(*mFinishedSum); - for (int32_t bi = 0; bi < localBatchSize; ++bi) - { - finishedSumHost[bi] = 0; - } - } - - if (isBeamSearch) - { - auto outputsBeamSearch = std::dynamic_pointer_cast(outputParams); - TLLM_CHECK_WITH_INFO(tgtCacheIndirectionOpt.has_value(), "tgtCacheIndirection must be set for beam search"); - outputsBeamSearch->tgtCacheIndirection = convert_tensor(tgtCacheIndirectionOpt.value()); - if (useBeamHyps) - { - // Additional parameters for beam search - outputsBeamSearch->beamHypotheses = std::make_unique(); - safeUpdatePtr(beamHypsIsDoneOpt, outputsBeamSearch->beamHypotheses->batchDones); - safeUpdatePtr(beamHypsCumLogProbsCbaOpt, outputsBeamSearch->beamHypotheses->cumLogProbsCBA); - safeUpdatePtr(beamHypsLogProbsCbaOpt, outputsBeamSearch->beamHypotheses->logProbsCBA); - safeUpdatePtr(beamHypsMinNormedScoresOpt, outputsBeamSearch->beamHypotheses->minNormedScoresCBA); - safeUpdatePtr(beamHypsNormedScoresCbaOpt, outputsBeamSearch->beamHypotheses->normedScoresCBA); - safeUpdatePtr(beamHypsNumBeamsOpt, outputsBeamSearch->beamHypotheses->numBeamsCBA); - safeUpdatePtr(beamHypsOutputIdsCbaOpt, outputsBeamSearch->beamHypotheses->outputIdsCBA); - safeUpdatePtr(beamHypsSeqLenCbaOpt, outputsBeamSearch->beamHypotheses->sequenceLengthsCBA); - } - } - - mDynamicDecodeLayer->forwardAsync(outputParams, forwardParams, mDecodingWorkspace); - - if (finishedSumHost) - { - TLLM_CUDA_CHECK(::cudaStreamSynchronize(mDynamicDecodeLayer->getStream())); - uint32_t numRealFinished = 0; - for (int32_t bi = 0; bi < localBatchSize; ++bi) - { - numRealFinished += finishedSumHost[bi]; - } - auto const numToFinish = outputParams->finished.value()->getSize(); - auto shouldStopAccessor = shouldStop.accessor(); - shouldStopAccessor[0] = numToFinish == numRealFinished; - } -} - -DynamicDecodeOp::DynamicDecodeOp(int64_t const maxBatchSize, int64_t const maxBeamWidth, int64_t const vocabSize, - int64_t const vocabSizePadded, int64_t const tensorParaSize, int64_t const pipelineParaSize, - at::ScalarType const scalarType) - : maxBatchSize_(static_cast(maxBatchSize)) - , maxBeamWidth_(static_cast(maxBeamWidth)) - , vocabSize_(static_cast(vocabSize)) - , vocabSizePadded_(static_cast(vocabSizePadded)) - , tensorParaSize_(static_cast(tensorParaSize)) - , pipelineParaSize_(static_cast(pipelineParaSize)) - , scalarType_(scalarType) -{ - TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); - createInstance(); -} - -void DynamicDecodeOp::createInstance() -{ - dynamicDecode_.reset(); - switch (scalarType_) - { - case at::ScalarType::Float: - dynamicDecode_ = std::make_unique>( - maxBatchSize_, maxBeamWidth_, vocabSize_, vocabSizePadded_, tensorParaSize_, pipelineParaSize_); - break; - case at::ScalarType::Half: - dynamicDecode_ = std::make_unique>( - maxBatchSize_, maxBeamWidth_, vocabSize_, vocabSizePadded_, tensorParaSize_, pipelineParaSize_); - break; - default: throw std::runtime_error("Wrong tensor type."); - } -} - -void DynamicDecodeOp::setup(int64_t const batchSize, int64_t const beamWidth, th::optional runtimeTopKOpt, - th::optional runtimeTopPOpt, th::optional temperatureOpt, - th::optional repetitionPenaltyOpt, th::optional presencePenaltyOpt, - th::optional frequencyPenaltyOpt, th::optional promptIgnoreLengthOpt, - th::optional minLengthOpt, th::optional lengthPenaltyOpt, - th::optional earlyStoppingOpt, th::optional beamSearchDiversityRateOpt, - th::optional randomSeedOpt, th::optional topPDecayOpt, th::optional topPMinOpt, - th::optional topPResetIdsOpt, th::optional noRepeatNgramSizeOpt, - th::optional minPOpt, bool outputLogProbs, bool cumLogProbs) -{ - // TODO: Revise DynamicDecodeLayer and make the decode arguments consistent. - // TODO: add parameters "normalizeLogProbs" and "topKMedusaHeads" - CHECK_OPTIONAL_CPU_INPUT(runtimeTopKOpt, torch::kInt32); - CHECK_OPTIONAL_CPU_INPUT(runtimeTopPOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(temperatureOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(repetitionPenaltyOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(presencePenaltyOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(frequencyPenaltyOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(promptIgnoreLengthOpt, torch::kInt32); - CHECK_OPTIONAL_CPU_INPUT(minLengthOpt, torch::kInt32); - CHECK_OPTIONAL_CPU_INPUT(lengthPenaltyOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(earlyStoppingOpt, torch::kInt32); - CHECK_OPTIONAL_CPU_INPUT(beamSearchDiversityRateOpt, torch::kFloat); - CHECK_OPTIONAL_CPU_INPUT(randomSeedOpt, torch::kInt64); - CHECK_OPTIONAL_INPUT(topPDecayOpt, torch::kFloat); - CHECK_OPTIONAL_INPUT(topPMinOpt, torch::kFloat); - CHECK_OPTIONAL_INPUT(topPResetIdsOpt, torch::kInt32); - CHECK_OPTIONAL_CPU_INPUT(noRepeatNgramSizeOpt, torch::kInt32); - CHECK_OPTIONAL_CPU_INPUT(minPOpt, torch::kFloat); - - dynamicDecode_->setup(static_cast(batchSize), static_cast(beamWidth), - runtimeTopKOpt, runtimeTopPOpt, temperatureOpt, repetitionPenaltyOpt, presencePenaltyOpt, frequencyPenaltyOpt, - promptIgnoreLengthOpt, minLengthOpt, lengthPenaltyOpt, earlyStoppingOpt, beamSearchDiversityRateOpt, - randomSeedOpt, topPDecayOpt, topPMinOpt, topPResetIdsOpt, noRepeatNgramSizeOpt, minPOpt, outputLogProbs, - cumLogProbs); -} - -th::Tensor DynamicDecodeOp::forward( - // Inputs BS: batchSize, BM: beamWidth, MSL: maxSeqLength, V: vocabSize, VP: vocabSizePadded - th::Tensor const& logits, // [BS, BM, VP], T, variables for input - int64_t const step, // - int64_t const maxInputLength, // - int64_t const maxAttentionWindow, // - int64_t const sinkTokenLength, // - int64_t const ite, // - int64_t const localBatchSize, // - th::Tensor const endId, // [BS*BM], int - th::optional embeddingBiasOpt, // [VP], T - th::optional inputLengthsOpt, // [BS*BM], int, length of input contexts - th::optional sequenceLimitLengthOpt, // [BS, 1], int - th::optional stopWordsListPtrsOpt, // [BS][2, stopWordsLength], int64 - th::optional stopWordsLensOpt, // [BS], int - int64_t const maxStopWordsLen, // - th::optional badWordsListPtrsOpt, // [BS][2, badWordsLength], int64 - th::optional badWordsLensOpt, // [BS], int - int64_t const maxBadWordsLen, // - th::optional srcCacheIndirectionOpt, // [localBS, BM, MSL], int - // Outputs - th::Tensor outputTokenIds, // [BS, BM, MSL], variables for output - th::Tensor newTokens, // [BS, BM, 1], int - th::optional finishedInput, // [BS, BM], uint8 - th::optional finishedOutput, // [BS, BM], uint8 - th::optional sequenceLengthsOpt, // [BS*BM], int, length of the current sequences - th::optional cumLogProbsOpt, // [BS, BM], float - th::optional outputLogProbsOpt, // [BS, BM, MSL], float - th::optional outputLogProbsTiledOpt, // [MSL, BS, BM], float, transpose of outputLogProbsOpt - th::optional parentIdsOpt, // [BS, BM, MSL], int - th::optional tgtCacheIndirectionOpt, // [localBS, BM, MSL], int - th::optional beamHypsOutputIdsCbaOpt, // [BS, BM*2, MSL], int - th::optional beamHypsSeqLenCbaOpt, // [BS, BM*2], int - th::optional beamHypsCumLogProbsCbaOpt, // [BS, BM*2], float - th::optional beamHypsNormedScoresCbaOpt, // [BS, BM*2], float - th::optional beamHypsLogProbsCbaOpt, // [BS, BM*2, MSL], float - th::optional beamHypsMinNormedScoresOpt, // [BS], float - th::optional beamHypsNumBeamsOpt, // [BS], int - th::optional beamHypsIsDoneOpt, // [BS], bool - bool const useBeamHyps // -) -{ - CHECK_INPUT(logits, scalarType_); - TLLM_CHECK_WITH_INFO(logits.dim() == 3, - "logits is of shape (batchSize, beamWidth, vocabSizePadded), but got dim=%d shape=%s", (int) logits.dim(), - tensorrt_llm::runtime::ITensor::toString(convert_shape(logits)).c_str()); - TLLM_CHECK_WITH_INFO(static_cast(logits.size(2)) == vocabSizePadded_, - "logits is of shape (batchSize, beamWidth, vocabSize(%ld)), but got the last dim=%ld.", vocabSizePadded_, - static_cast(logits.size(2))); - CHECK_INPUT(endId, torch::kInt32); - CHECK_OPTIONAL_INPUT(embeddingBiasOpt, scalarType_); - CHECK_OPTIONAL_INPUT(inputLengthsOpt, torch::kInt32); - CHECK_OPTIONAL_INPUT(sequenceLimitLengthOpt, torch::kInt32); - CHECK_OPTIONAL_INPUT(stopWordsListPtrsOpt, torch::kInt64); - CHECK_OPTIONAL_INPUT(stopWordsLensOpt, torch::kInt32); - CHECK_OPTIONAL_INPUT(badWordsListPtrsOpt, torch::kInt64); - CHECK_OPTIONAL_INPUT(badWordsLensOpt, torch::kInt32); - CHECK_OPTIONAL_INPUT(srcCacheIndirectionOpt, torch::kInt32); - CHECK_INPUT(outputTokenIds, torch::kInt32); - CHECK_INPUT(newTokens, torch::kInt32); - CHECK_OPTIONAL_INPUT(finishedInput, torch::kUInt8); - CHECK_OPTIONAL_INPUT(finishedOutput, torch::kUInt8); - CHECK_OPTIONAL_INPUT(sequenceLengthsOpt, torch::kInt32); - CHECK_OPTIONAL_INPUT(cumLogProbsOpt, torch::kFloat32); - CHECK_OPTIONAL_INPUT(outputLogProbsOpt, torch::kFloat32); - CHECK_OPTIONAL_INPUT(outputLogProbsTiledOpt, torch::kFloat32); - CHECK_OPTIONAL_INPUT(parentIdsOpt, torch::kInt32); - CHECK_OPTIONAL_INPUT(tgtCacheIndirectionOpt, torch::kInt32); - - th::Tensor shouldStop = torch::zeros({1}, torch::dtype(torch::kBool).requires_grad(false)); - - dynamicDecode_->forward( - // Inputs - logits, static_cast(step), static_cast(maxInputLength), static_cast(maxAttentionWindow), - static_cast(sinkTokenLength), static_cast(ite), static_cast(localBatchSize), endId, - embeddingBiasOpt, inputLengthsOpt, sequenceLimitLengthOpt, stopWordsListPtrsOpt, stopWordsLensOpt, - static_cast(maxStopWordsLen), badWordsListPtrsOpt, badWordsLensOpt, - static_cast(maxBadWordsLen), srcCacheIndirectionOpt, - // Outputs - outputTokenIds, newTokens, shouldStop, finishedInput, finishedOutput, sequenceLengthsOpt, cumLogProbsOpt, - outputLogProbsOpt, outputLogProbsTiledOpt, parentIdsOpt, tgtCacheIndirectionOpt, beamHypsOutputIdsCbaOpt, - beamHypsSeqLenCbaOpt, beamHypsCumLogProbsCbaOpt, beamHypsNormedScoresCbaOpt, beamHypsLogProbsCbaOpt, - beamHypsMinNormedScoresOpt, beamHypsNumBeamsOpt, beamHypsIsDoneOpt, useBeamHyps); - - return shouldStop; -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -static auto trtllmGptContextDecoderTHS - = torch::jit::class_("trtllm", "DynamicDecodeOp") - .def(torch::jit::init()) - .def("setup", &tensorrt_llm::torch_ext::DynamicDecodeOp::setup) - .def("forward", &tensorrt_llm::torch_ext::DynamicDecodeOp::forward); diff --git a/cpp/tensorrt_llm/thop/dynamicDecodeOp.h b/cpp/tensorrt_llm/thop/dynamicDecodeOp.h deleted file mode 100644 index c8f4fa807d31..000000000000 --- a/cpp/tensorrt_llm/thop/dynamicDecodeOp.h +++ /dev/null @@ -1,165 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaBf16Wrapper.h" -#include "tensorrt_llm/layers/dynamicDecodeLayer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/thop/thUtils.h" - -namespace th = torch; - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -class IFtDynamicDecode -{ -public: - virtual ~IFtDynamicDecode() = default; - - virtual void setup(size_t const batch_size, size_t const beam_width, th::optional runtime_top_k_opt, - th::optional runtime_top_p_opt, th::optional temperature_opt, - th::optional repetition_penalty_opt, th::optional presence_penalty_opt, - th::optional frequency_penalty_opt, th::optional prompt_ignore_length_opt, - th::optional min_length_opt, th::optional length_penalty_opt, - th::optional early_stopping_opt, th::optional beam_search_diversity_rate_opt, - th::optional random_seed_opt, th::optional top_p_decay_opt, - th::optional top_p_min_opt, th::optional top_p_reset_ids_opt, - th::optional no_repeat_ngram_size_opt, th::optional min_p_opt, bool output_log_probs, - bool cum_log_probs) - = 0; - - virtual void forward(th::Tensor const& logits, int const step, int const max_input_length, - int const max_attention_window, int const sink_token_length, uint64_t const ite, int const local_batch_size, - th::Tensor end_id, th::optional embedding_bias_opt, th::optional input_lengths_opt, - th::optional sequence_limit_length_opt, th::optional stop_words_list_ptrs_opt, - th::optional stop_words_lens_opt, int32_t const max_stop_words_len, - th::optional bad_words_list_ptrs_opt, th::optional bad_words_lens_opt, - int32_t const max_bad_words_len, th::optional src_cache_indirection_opt, - th::Tensor& output_token_ids, th::Tensor& newTokens, th::Tensor& should_stop, - th::optional finished_input, th::optional finished_output, - th::optional sequence_lengths_opt, th::optional cum_log_probs_opt, - th::optional output_log_probs_opt, th::optional output_log_probs_tiled_opt, - th::optional parent_ids_opt, th::optional tgt_cache_indirection_opt, - th::optional beam_hyps_output_ids_cba_opt, th::optional beam_hyps_seq_len_cba_opt, - th::optional beam_hyps_cum_log_probs_cba_opt, - th::optional beam_hyps_normed_scores_cba_opt, th::optional beam_hyps_log_probs_cba_opt, - th::optional beam_hyps_min_normed_scores_opt, th::optional beam_hyps_num_beams_opt, - th::optional beam_hyps_is_done_opt, bool const use_beam_hyps) - = 0; -}; - -template -class FtDynamicDecode : public IFtDynamicDecode -{ -public: - FtDynamicDecode(size_t const max_batch_size, size_t const max_beam_width, size_t const vocab_size, - size_t const vocab_size_padded, int const tensor_para_size, int const pipeline_para_size); - - ~FtDynamicDecode() override = default; - - void setup(size_t const batch_size, size_t const beam_width, th::optional runtime_top_k_opt, - th::optional runtime_top_p_opt, th::optional temperature_opt, - th::optional repetition_penalty_opt, th::optional presence_penalty_opt, - th::optional frequency_penalty_opt, th::optional prompt_ignore_length_opt, - th::optional min_length_opt, th::optional length_penalty_opt, - th::optional early_stopping_opt, th::optional beam_search_diversity_rate_opt, - th::optional random_seed_opt, th::optional top_p_decay_opt, - th::optional top_p_min_opt, th::optional top_p_reset_ids_opt, - th::optional no_repeat_ngram_size_opt, th::optional min_p_opt, bool output_log_probs, - bool cum_log_probs) override; - - void forward(th::Tensor const& logits, int const step, int const max_input_length, int const max_attention_window, - int const sink_token_length, uint64_t const ite, int const local_batch_size, th::Tensor end_id, - th::optional embedding_bias_opt, th::optional input_lengths_opt, - th::optional sequence_limit_length_opt, th::optional stop_words_list_ptrs_opt, - th::optional stop_words_lens_opt, int32_t const max_stop_words_len, - th::optional bad_words_list_ptrs_opt, th::optional bad_words_lens_opt, - int32_t const max_bad_words_len, th::optional src_cache_indirection_opt, - th::Tensor& output_token_ids, th::Tensor& newTokens, th::Tensor& should_stop, - th::optional finished_input, th::optional finished_output, - th::optional sequence_lengths_opt, th::optional cum_log_probs_opt, - th::optional output_log_probs_opt, th::optional output_log_probs_tiled_opt, - th::optional parent_ids_opt, th::optional tgt_cache_indirection_opt, - th::optional beam_hyps_output_ids_cba_opt, th::optional beam_hyps_seq_len_cba_opt, - th::optional beam_hyps_cum_log_probs_cba_opt, - th::optional beam_hyps_normed_scores_cba_opt, th::optional beam_hyps_log_probs_cba_opt, - th::optional beam_hyps_min_normed_scores_opt, th::optional beam_hyps_num_beams_opt, - th::optional beam_hyps_is_done_opt, bool const use_beam_hyps) override; - -private: - tensorrt_llm::runtime::ITensor::SharedPtr mFinishedSum; // [batch_size] pinned - std::shared_ptr> mDynamicDecodeLayer; - std::shared_ptr mDecodingWorkspace; - std::optional mBeamWidth; - tensorrt_llm::runtime::ITensor::SharedConstPtr mBatchSlots; -}; - -class DynamicDecodeOp : public th::jit::CustomClassHolder -{ -public: - DynamicDecodeOp(int64_t const max_batch_size, int64_t const max_beam_width, int64_t const vocab_size, - int64_t const vocab_size_padded, int64_t const tensor_para_size, int64_t const pipeline_para_size, - at::ScalarType const scalar_type); - - void setup(int64_t const batch_size, int64_t const beam_width, th::optional runtime_top_k_opt, - th::optional runtime_top_p_opt, th::optional temperature_opt, - th::optional repetition_penalty_opt, th::optional presence_penalty_opt, - th::optional frequency_penalty_opt, th::optional prompt_ignore_length_opt, - th::optional min_length_opt, th::optional length_penalty_opt, - th::optional early_stopping_opt, th::optional beam_search_diversity_rate_opt, - th::optional random_seed_opt, th::optional top_p_decay_opt, - th::optional top_p_min_opt, th::optional top_p_reset_ids_opt, - th::optional no_repeat_ngram_size_opt, th::optional min_p_opt, bool output_log_probs, - bool cum_log_probs); - - th::Tensor forward(th::Tensor const& logits, int64_t const step, int64_t const max_input_length, - int64_t const max_attention_window, int64_t const sink_token_length, int64_t const ite, - int64_t const local_batch_size, th::Tensor end_id, th::optional embedding_bias_opt, - th::optional input_lengths_opt, th::optional sequence_limit_length_opt, - th::optional stop_words_list_ptrs_opt, th::optional stop_words_lens_opt, - int64_t const max_stop_words_len, th::optional bad_words_list_ptrs_opt, - th::optional bad_words_lens_opt, int64_t const max_bad_words_len, - th::optional src_cache_indirection_opt, th::Tensor output_token_ids, th::Tensor newTokens, - th::optional finished_input, th::optional finished_output, - th::optional sequence_lengths_opt, th::optional cum_log_probs_opt, - th::optional output_log_probs_opt, th::optional output_log_probs_tiled_opt, - th::optional parent_ids_opt, th::optional tgt_cache_indirection_opt, - th::optional beam_hyps_output_ids_cba_opt, th::optional beam_hyps_seq_len_cba_opt, - th::optional beam_hyps_cum_log_probs_cba_opt, - th::optional beam_hyps_normed_scores_cba_opt, th::optional beam_hyps_log_probs_cba_opt, - th::optional beam_hyps_min_normed_scores_opt, th::optional beam_hyps_num_beams_opt, - th::optional beam_hyps_is_done_opt, bool const use_beam_hyps); - -private: - // Members initialized in constructor and used in call of createInstance() - size_t const maxBatchSize_; - size_t const maxBeamWidth_; - size_t const vocabSize_; - size_t const vocabSizePadded_; - int const tensorParaSize_; - int const pipelineParaSize_; - at::ScalarType const scalarType_; // Data type of expected input logits - std::unique_ptr dynamicDecode_; // FT Dynamic decode layer wrapper instance - - void createInstance(); -}; - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/gatherTreeOp.cpp b/cpp/tensorrt_llm/thop/gatherTreeOp.cpp deleted file mode 100644 index 45f2649a6a71..000000000000 --- a/cpp/tensorrt_llm/thop/gatherTreeOp.cpp +++ /dev/null @@ -1,142 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/decodingKernels.h" -#include "tensorrt_llm/thop/thUtils.h" - -namespace th = torch; -namespace tl = tensorrt_llm; -namespace tk = tensorrt_llm::kernels; - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -// Must be similar to GptDecoder::gatherTree -th::Tensor gatherTree( // BS: batch_size, BM: beam_width, MSL: max_seq_length - th::Tensor& sequence_lengths, // [BS*BM], int - th::Tensor& output_ids, // [BS, BM, MSL],int - th::Tensor& parent_ids, // [BS, BM, MSL], int - th::Tensor& end_ids, // [BS*BM], int - th::Tensor& tiled_input_lengths, // [BS*BM], int - th::optional cum_log_probs_opt, // [BS, BM], float - th::optional log_probs_opt, // [BS, BM, MSL], float - th::optional log_probs_tiled_opt, // [MSL, BS, BM], float, transpose of output_log_probs_opt - th::optional beam_hyps_output_ids_cba, // [BS, BM*2, MSL], int - th::optional beam_hyps_seq_len_cba, // [BS, BM*2], int - th::optional beam_hyps_cum_log_probs_cba, // [BS, BM*2], float - th::optional beam_hyps_normed_scores_cba, // [BS, BM*2], float - th::optional beam_hyps_log_probs_cba, // [BS, BM*2, MSL], float - th::optional beam_hyps_min_normed_scores, // [BS], float - th::optional beam_hyps_num_beams, // [BS], int - th::optional beam_hyps_is_done, // [BS], bool - th::optional finished, // [BS, BM], uint8 - th::Tensor& length_penalty, // [BS], float - int64_t const batch_size, // - int64_t const beam_width, // - int64_t const max_seq_len, // - bool const use_beam_hyps // -) -{ - auto stream = at::cuda::getCurrentCUDAStream().stream(); - th::Tensor final_output_ids = torch::zeros( - {batch_size, beam_width, max_seq_len}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); - if (use_beam_hyps && beam_width > 1) - { - int32_t* final_output_ids_ptr = get_ptr(final_output_ids); - tk::invokeInitializeOutput( - final_output_ids_ptr, get_ptr(end_ids), batch_size, beam_width, max_seq_len, stream); - - tk::BeamHypotheses bh; - bh.nBatchSize = batch_size; - bh.nBeamWidth = beam_width; - bh.nMaxSeqLen = max_seq_len; - bh.lengthPenalties = get_ptr(length_penalty); - bh.inputLengths = get_ptr(tiled_input_lengths); - bh.outputIds = final_output_ids_ptr; - bh.logProbs = log_probs_opt.has_value() ? get_ptr(log_probs_opt.value()) : nullptr; - bh.logProbsTiled = log_probs_tiled_opt.has_value() ? get_ptr(log_probs_tiled_opt.value()) : nullptr; - bh.sequenceLengths = get_ptr(sequence_lengths); - bh.cumLogProbs = cum_log_probs_opt.has_value() ? get_ptr(cum_log_probs_opt.value()) : nullptr; - bh.outputIdsCBA = get_ptr(beam_hyps_output_ids_cba.value()); - bh.logProbsCBA = get_ptr(beam_hyps_log_probs_cba.value()); - bh.sequenceLengthsCBA = get_ptr(beam_hyps_seq_len_cba.value()); - bh.cumLogProbsCBA = get_ptr(beam_hyps_cum_log_probs_cba.value()); - bh.normedScoresCBA = get_ptr(beam_hyps_normed_scores_cba.value()); - bh.numBeamsCBA = get_ptr(beam_hyps_num_beams.value()); - bh.minNormedScoresCBA = get_ptr(beam_hyps_min_normed_scores.value()); - bh.batchDones = get_ptr(beam_hyps_is_done.value()); - bh.finished - = reinterpret_cast(get_ptr(finished.value())); - bh.outputIdsUnfinish = get_ptr(output_ids); - bh.parentIdsUnfinish = get_ptr(parent_ids); - - tk::invokeInsertUnfinishedPath(bh, stream); - sync_check_cuda_error(stream); - - tk::invokeFinalize(bh, stream); - sync_check_cuda_error(stream); - } - else if (!use_beam_hyps && beam_width > 1) - { - th::Tensor workspace = torch::zeros(batch_size * beam_width * max_seq_len * sizeof(int32_t), - torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false)); - - // For sampling, it is equivalent to all parent ids are 0. - tk::gatherTreeParam param; - param.beams = get_ptr(workspace); - // Remove prompt length if possible - param.sequenceLengths = get_ptr(sequence_lengths); - // add sequence_length 1 here because the sequence_length of time step t is t - 1 - param.maxSequenceLengthFinalStep = 1; - // response input lengths (used to slice the ids during postprocessing), used in interactive generation - // This feature is not supported yet, setting it to nullptr temporarily. - param.responseInputLengths = nullptr; - param.maxSeqLen = max_seq_len; - param.batchSize = batch_size; - param.beamWidth = beam_width; - param.stepIds = get_ptr(output_ids); - param.parentIds = beam_width == 1 ? nullptr : get_ptr(parent_ids); - param.endTokens = get_ptr(end_ids); - param.inputLengths = get_ptr(tiled_input_lengths); - - param.stream = stream; - param.outputIds = get_ptr(final_output_ids); - param.cumLogProbs = cum_log_probs_opt.has_value() ? get_ptr(cum_log_probs_opt.value()) : nullptr; - param.lengthPenalty = get_val(length_penalty, 0); - - // NOTE: need to remove all prompt virtual tokens - tk::invokeGatherTree(param); - sync_check_cuda_error(stream); - } - else - { - cudaMemcpyAsync(get_ptr(final_output_ids), get_ptr(output_ids), - sizeof(int) * batch_size * beam_width * max_seq_len, cudaMemcpyDeviceToDevice, stream); - sync_check_cuda_error(stream); - } - return final_output_ids; -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -static auto gather_tree = torch::RegisterOperators("tensorrt_llm::gather_tree", &tensorrt_llm::torch_ext::gatherTree); diff --git a/cpp/tests/unit_tests/CMakeLists.txt b/cpp/tests/unit_tests/CMakeLists.txt index 9d22bd03b52a..8db945f04101 100644 --- a/cpp/tests/unit_tests/CMakeLists.txt +++ b/cpp/tests/unit_tests/CMakeLists.txt @@ -24,6 +24,5 @@ endif() add_subdirectory(common) add_subdirectory(kernels) add_subdirectory(multi_gpu) -add_subdirectory(layers) add_subdirectory(runtime) add_subdirectory(thop) diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 6b0e5a118211..9b534909fd4d 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -14,7 +14,6 @@ # the License. add_gtest(banRepeatNGramsKernelsTest banRepeatNGramsKernelsTest.cpp) -add_gtest(decodingKernelsTest decodingKernelTest.cpp) add_gtest(logitsBitmaskTest logitsBitmaskTest.cpp) add_gtest(cascadeAttentionKernelTest cascadeAttentionKernelTest.cpp) add_gtest(cascadeAttentionNumericsTest cascadeAttentionNumericsTest.cu) @@ -110,6 +109,5 @@ if(USING_OSS_CUTLASS_MOE_GEMM) add_gtest(moeLoraSlotExpandTest moeLoraSlotExpandTest.cu) endif() -add_gtest(eaglePackDataTest eaglePackDataTest.cpp) add_gtest(sparseKvCacheTest sparseKvCacheTest.cu) add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp) diff --git a/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp b/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp deleted file mode 100644 index a458f783ef65..000000000000 --- a/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp +++ /dev/null @@ -1,2032 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/decodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -#include -#include - -#include -#include -#include -#include - -namespace tk = tensorrt_llm::kernels; -namespace tksp = tensorrt_llm::kernels::speculative_decoding; -namespace tc = tensorrt_llm::common; -namespace trk = tensorrt_llm::runtime::kernels; - -using namespace tensorrt_llm::runtime; - -namespace -{ - -inline bool almostEqual(float a, float b, float atol = 1e-2, float rtol = 1e-3) -{ - // Params: a = value to compare and b = reference - // This function follows implementation of numpy.isclose(), which checks - // abs(a - b) <= (atol + rtol * abs(b)). - // Note that the inequality above is asymmetric where b is considered as - // a reference value. To account into both absolute/relative errors, it - // uses absolute tolerance and relative tolerance at the same time. The - // default values of atol and rtol borrowed from numpy.isclose(). For the - // case of nan value, the result will be true. - if (isnan(a) && isnan(b)) - { - return true; - } - return fabs(a - b) <= (atol + rtol * fabs(b)); -} - -std::vector calculateGaussianKernel(float sigma, int size) -{ - std::vector kernel(size); - float sum = 0.f; - - for (int i = 0; i < size; ++i) - { - int x = i - size / 2; - kernel[i] = std::exp(-0.5f * (x * x) / (sigma * sigma)); - sum += kernel[i]; - } - - // Normalize the kernel - for (int i = 0; i < size; ++i) - { - kernel[i] /= sum; - } - - return kernel; -} - -template -void applyGaussianFilter(T* result, float const* input, int n, float sigma) -{ - int size = static_cast(std::ceil(6.f * sigma)); - size = (size % 2 == 0) ? size + 1 : size; - - std::vector kernel = calculateGaussianKernel(sigma, size); - int halfSize = size / 2; - - for (int i = 0; i < n; ++i) - { - result[i] = T{0}; - } - - // Convolution operation - for (int i = 0; i < n; ++i) - { - for (int j = 0; j < size; ++j) - { - int k = i - halfSize + j; - if (k >= 0 && k < n) - { - result[i] += input[k] * kernel[j]; - } - } - } -} - -template void applyGaussianFilter(float* result, float const* input, int n, float sigma); -template void applyGaussianFilter(__half* result, float const* input, int n, float sigma); - -template -void probsToLogits(T const* probs, T* logits, SizeType32 n) -{ - constexpr float eps = 1e-6f; - for (SizeType32 ni = 0; ni < n; ++ni) - { - auto const prob = std::max(eps, static_cast(probs[ni])); - logits[ni] = std::log(prob / (1.f - prob)); - } -} - -template -void softmax(T const* logits, T* probs, int n) -{ - float epsilon = 1e-6f; - - // Find the maximum logit value - float maxLogits = -std::numeric_limits::max(); - for (int ii = 0; ii < n; ++ii) - { - maxLogits = std::max(maxLogits, static_cast(logits[ii])); - } - - // Calculate the numerator of the softmax formula - float expSum = 0.0; - for (int ii = 0; ii < n; ++ii) - { - expSum += std::exp(static_cast(logits[ii]) - maxLogits); - } - - // Calculate softmax probabilities - for (int ii = 0; ii < n; ++ii) - { - float prob = std::exp(static_cast(logits[ii]) - maxLogits) / (expSum + epsilon); - probs[ii] = prob; - } -} - -template void probsToLogits(float const* probs, float* logits, SizeType32 n); -template void probsToLogits(__half const* probs, __half* logits, SizeType32 n); - -template -void checkEquality(DecodingOutput::TensorPtr src, DecodingOutput::TensorPtr dst, char const* bufferName, - tensorrt_llm::runtime::BufferManager& bufferManager) -{ - auto srcHost = bufferManager.copyFrom(*src, MemoryType::kPINNEDPOOL); - auto dstHost = bufferManager.copyFrom(*dst, MemoryType::kPINNEDPOOL); - bufferManager.getStream().synchronize(); - auto srcPtr = bufferCast(*srcHost); - auto dstPtr = bufferCast(*dstHost); - for (SizeType32 ii = 0; ii < src->getSize(); ++ii) - { - // since it's a simple copy, floats support the simple equality - EXPECT_EQ(srcPtr[ii], dstPtr[ii]) << "Unequal values in buffer " << bufferName << " at ii: " << ii - << " with values: src " << srcPtr[ii] << " dst " << dstPtr[ii] << std::endl; - } -} - -template -void fillBufferWithRandom(ITensor& buffer, tensorrt_llm::runtime::BufferManager& bufferManager, std::mt19937& randGen) -{ - auto cpuBuffer = bufferManager.cpu(buffer.getShape(), TRTDataType::value); - - auto const size = cpuBuffer->getSize(); - auto rawPtr = bufferCast(*cpuBuffer); - - std::uniform_int_distribution<> dis(0, 255); - - for (SizeType32 i = 0; i < size; ++i) - { - rawPtr[i] = static_cast(dis(randGen)); - } - bufferManager.copy(*cpuBuffer, buffer); -} - -class TestBeamHypothesesCopy : public ::testing::Test -{ -public: - DecodingOutput::BeamHypotheses srcBeams; - DecodingOutput::BeamHypotheses dstBeams; - DecodingOutput::TensorPtr mSrcCumLogProbs; - DecodingOutput::TensorPtr mDstCumLogProbs; - SizeType32 mNumSMs; - - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - - std::mt19937 gen; - - void SetUp() override - { - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); - int device; - cudaGetDevice(&device); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, device); - mNumSMs = deviceProp.multiProcessorCount; - gen.seed(42U); - } - - void initializeBuffers(SizeType32 batchSize, SizeType32 beamWidth, SizeType32 maxSeqLen) - { - - srcBeams.empty(*mBufferManager); - srcBeams.reshape(batchSize, beamWidth, maxSeqLen); - mSrcCumLogProbs - = mBufferManager->gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kFLOAT); - - setBuffers(srcBeams, mSrcCumLogProbs, 2); - - dstBeams.empty(*mBufferManager); - dstBeams.reshape(batchSize, beamWidth, maxSeqLen); - mDstCumLogProbs - = mBufferManager->gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kFLOAT); - - setBuffers(dstBeams, mDstCumLogProbs, 1); - } - - void setBuffers(DecodingOutput::BeamHypotheses currBeams, DecodingOutput::TensorPtr cumLogProbs, int value) - { - fillBufferWithRandom(*currBeams.outputIdsCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.logProbsCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.sequenceLengthsCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.cumLogProbsCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.normedScoresCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.numBeamsCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.minNormedScoresCBA, *mBufferManager, gen); - fillBufferWithRandom(*currBeams.batchDones, *mBufferManager, gen); - fillBufferWithRandom(*cumLogProbs, *mBufferManager, gen); - } - - void checkAllEqual() - { - checkEquality(srcBeams.outputIdsCBA, dstBeams.outputIdsCBA, "outputIdsCBA", *mBufferManager); - checkEquality(srcBeams.logProbsCBA, dstBeams.logProbsCBA, "logProbsCBA", *mBufferManager); - checkEquality( - srcBeams.sequenceLengthsCBA, dstBeams.sequenceLengthsCBA, "sequenceLengthsCBA", *mBufferManager); - checkEquality(srcBeams.cumLogProbsCBA, dstBeams.cumLogProbsCBA, "cumLogProbsCBA", *mBufferManager); - checkEquality(srcBeams.normedScoresCBA, dstBeams.normedScoresCBA, "normedScoresCBA", *mBufferManager); - checkEquality(srcBeams.numBeamsCBA, dstBeams.numBeamsCBA, "numBeamsCBA", *mBufferManager); - checkEquality( - srcBeams.minNormedScoresCBA, dstBeams.minNormedScoresCBA, "minNormedScoresCBA", *mBufferManager); - checkEquality(srcBeams.batchDones, dstBeams.batchDones, "batchDones", *mBufferManager); - checkEquality(mSrcCumLogProbs, mDstCumLogProbs, "cumLogProbs", *mBufferManager); - } -}; - -// Test for invokeCopyBeamHypotheses -TEST_F(TestBeamHypothesesCopy, FullBatchTest) -{ - SizeType32 const batchSize{1024}; - SizeType32 const beamWidth{64}; - SizeType32 const maxSeqLen{2048}; - - initializeBuffers(batchSize, beamWidth, maxSeqLen); - mStream->synchronize(); - - tk::invokeCopyBeamHypotheses(srcBeams, dstBeams, *mSrcCumLogProbs, *mDstCumLogProbs, *mStream, mNumSMs); - mStream->synchronize(); - - checkAllEqual(); -} - -TEST_F(TestBeamHypothesesCopy, SingleBatchTest) -{ - SizeType32 const batchSize{1}; - SizeType32 const beamWidth{64}; - SizeType32 const maxSeqLen{16384}; - - initializeBuffers(batchSize, beamWidth, maxSeqLen); - mStream->synchronize(); - - tk::invokeCopyBeamHypotheses(srcBeams, dstBeams, *mSrcCumLogProbs, *mDstCumLogProbs, *mStream, mNumSMs); - mStream->synchronize(); - - checkAllEqual(); -} - -void runFinalizeKeepsPromptPrefixFromUnfinishedIds(SizeType32 beamWidth) -{ - SizeType32 constexpr batchSize{2}; - SizeType32 constexpr maxSeqLen{9}; - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - auto constexpr nvFloatType = TRTDataType::value; - auto constexpr nvBoolType = TRTDataType::value; - - auto stream = std::make_shared(); - auto bufferManager = std::make_shared(stream); - - auto const idsShape = ITensor::makeShape({batchSize, beamWidth, maxSeqLen}); - auto const cbaShape = ITensor::makeShape({batchSize, beamWidth * 2, maxSeqLen}); - auto const beamShape = ITensor::makeShape({batchSize, beamWidth}); - auto const cbaBeamShape = ITensor::makeShape({batchSize, beamWidth * 2}); - - auto outputIds = bufferManager->gpu(idsShape, nvTokenIdType); - auto outputIdsUnfinish = bufferManager->gpu(idsShape, nvTokenIdType); - auto parentIdsUnfinish = bufferManager->gpu(idsShape, nvSizeType); - auto outputIdsCBA = bufferManager->gpu(cbaShape, nvTokenIdType); - auto sequenceLengths = bufferManager->gpu(beamShape, nvSizeType); - auto inputLengths = bufferManager->gpu(beamShape, nvSizeType); - auto cumLogProbs = bufferManager->gpu(beamShape, nvFloatType); - auto sequenceLengthsCBA = bufferManager->gpu(cbaBeamShape, nvSizeType); - auto cumLogProbsCBA = bufferManager->gpu(cbaBeamShape, nvFloatType); - auto normedScoresCBA = bufferManager->gpu(cbaBeamShape, nvFloatType); - auto numBeamsCBA = bufferManager->gpu(ITensor::makeShape({batchSize}), nvSizeType); - auto batchDones = bufferManager->gpu(ITensor::makeShape({batchSize}), nvBoolType); - auto lengthPenalties = bufferManager->gpu(ITensor::makeShape({batchSize}), nvFloatType); - - std::vector outputIdsHost(batchSize * beamWidth * maxSeqLen, 0); - std::vector outputIdsUnfinishHost(batchSize * beamWidth * maxSeqLen, -1); - std::vector parentIdsUnfinishHost(batchSize * beamWidth * maxSeqLen, 0); - std::vector outputIdsCBAHost(batchSize * beamWidth * 2 * maxSeqLen, -777); - std::vector sequenceLengthsHost(batchSize * beamWidth, 0); - std::vector inputLengthsHost(batchSize * beamWidth, 0); - std::vector cumLogProbsHost(batchSize * beamWidth, 0.0f); - std::vector sequenceLengthsCBAHost(batchSize * beamWidth * 2, 0); - std::vector cumLogProbsCBAHost(batchSize * beamWidth * 2, -100.0f); - std::vector normedScoresCBAHost(batchSize * beamWidth * 2, -100.0f); - std::vector numBeamsCBAHost(batchSize, beamWidth); - bool batchDonesHost[batchSize] = {false, false}; - float lengthPenaltiesHost[batchSize] = {0.0f, 0.0f}; - - auto promptToken = [](SizeType32 batchIdx, SizeType32 pos) -> TokenIdType - { return static_cast((batchIdx + 1) * 100 + pos); }; - auto unfinishedToken = [](SizeType32 batchIdx, SizeType32 beamIdx, SizeType32 pos) -> TokenIdType - { return static_cast(1000 + batchIdx * 100 + beamIdx * 10 + pos); }; - auto cbaToken = [](SizeType32 batchIdx, SizeType32 cbaIdx, SizeType32 pos) -> TokenIdType - { return static_cast(3000 + batchIdx * 100 + cbaIdx * 10 + pos); }; - auto inputLengthForBatch = [](SizeType32 batchIdx) -> SizeType32 { return batchIdx == 0 ? 3 : 5; }; - auto sequenceLengthForBatch = [](SizeType32 batchIdx) -> SizeType32 { return batchIdx == 0 ? 7 : 8; }; - - for (SizeType32 batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - SizeType32 const inputLength = inputLengthForBatch(batchIdx); - SizeType32 const sequenceLength = sequenceLengthForBatch(batchIdx); - SizeType32 const idsBatchOffset = batchIdx * beamWidth * maxSeqLen; - SizeType32 const cbaBatchOffset = batchIdx * beamWidth * 2 * maxSeqLen; - SizeType32 const cbaBeamOffset = batchIdx * beamWidth * 2; - - for (SizeType32 beamIdx = 0; beamIdx < beamWidth; ++beamIdx) - { - SizeType32 const beamOffset = idsBatchOffset + beamIdx * maxSeqLen; - inputLengthsHost[batchIdx * beamWidth + beamIdx] = inputLength; - sequenceLengthsHost[batchIdx * beamWidth + beamIdx] = sequenceLength; - cumLogProbsHost[batchIdx * beamWidth + beamIdx] - = beamWidth == 1 ? 12.0f : (beamIdx == 1 ? 12.0f : 5.0f + beamIdx * 3.0f); - - for (SizeType32 pos = 0; pos < sequenceLength; ++pos) - { - outputIdsUnfinishHost[beamOffset + pos] - = pos < inputLength ? promptToken(batchIdx, pos) : unfinishedToken(batchIdx, beamIdx, pos); - parentIdsUnfinishHost[beamOffset + pos] = beamIdx; - } - } - - for (SizeType32 cbaIdx = 0; cbaIdx < beamWidth; ++cbaIdx) - { - sequenceLengthsCBAHost[cbaBeamOffset + cbaIdx] = sequenceLength; - normedScoresCBAHost[cbaBeamOffset + cbaIdx] - = beamWidth == 1 ? 10.0f : (cbaIdx == 1 ? 10.0f : 1.0f - cbaIdx * 2.0f); - cumLogProbsCBAHost[cbaBeamOffset + cbaIdx] = normedScoresCBAHost[cbaBeamOffset + cbaIdx]; - for (SizeType32 pos = inputLength; pos < sequenceLength; ++pos) - { - outputIdsCBAHost[cbaBatchOffset + cbaIdx * maxSeqLen + pos] = cbaToken(batchIdx, cbaIdx, pos); - } - } - } - - bufferManager->copy(outputIdsHost.data(), *outputIds); - bufferManager->copy(outputIdsUnfinishHost.data(), *outputIdsUnfinish); - bufferManager->copy(parentIdsUnfinishHost.data(), *parentIdsUnfinish); - bufferManager->copy(outputIdsCBAHost.data(), *outputIdsCBA); - bufferManager->copy(sequenceLengthsHost.data(), *sequenceLengths); - bufferManager->copy(inputLengthsHost.data(), *inputLengths); - bufferManager->copy(cumLogProbsHost.data(), *cumLogProbs); - bufferManager->copy(sequenceLengthsCBAHost.data(), *sequenceLengthsCBA); - bufferManager->copy(cumLogProbsCBAHost.data(), *cumLogProbsCBA); - bufferManager->copy(normedScoresCBAHost.data(), *normedScoresCBA); - bufferManager->copy(numBeamsCBAHost.data(), *numBeamsCBA); - bufferManager->copy(batchDonesHost, *batchDones); - bufferManager->copy(lengthPenaltiesHost, *lengthPenalties); - stream->synchronize(); - - tk::BeamHypotheses bh; - bh.nMaxBatchSize = batchSize; - bh.nBatchSize = batchSize; - bh.nBeamWidth = beamWidth; - bh.nMaxSeqLen = maxSeqLen; - bh.lengthPenalties = bufferCast(*lengthPenalties); - bh.inputLengths = bufferCast(*inputLengths); - bh.outputIds = bufferCast(*outputIds); - bh.sequenceLengths = bufferCast(*sequenceLengths); - bh.cumLogProbs = bufferCast(*cumLogProbs); - bh.outputIdsCBA = bufferCast(*outputIdsCBA); - bh.sequenceLengthsCBA = bufferCast(*sequenceLengthsCBA); - bh.cumLogProbsCBA = bufferCast(*cumLogProbsCBA); - bh.normedScoresCBA = bufferCast(*normedScoresCBA); - bh.numBeamsCBA = bufferCast(*numBeamsCBA); - bh.batchDones = bufferCast(*batchDones); - bh.outputIdsUnfinish = bufferCast(*outputIdsUnfinish); - bh.parentIdsUnfinish = bufferCast(*parentIdsUnfinish); - - tk::invokeInsertUnfinishedPath(bh, stream->get()); - tk::invokeFinalize(bh, stream->get()); - stream->synchronize(); - - auto outputIdsResult = bufferManager->copyFrom(*outputIds, MemoryType::kCPU); - auto sequenceLengthsResult = bufferManager->copyFrom(*sequenceLengths, MemoryType::kCPU); - stream->synchronize(); - auto const outputIdsPtr = bufferCast(*outputIdsResult); - auto const sequenceLengthsPtr = bufferCast(*sequenceLengthsResult); - - for (SizeType32 batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - SizeType32 const inputLength = inputLengthForBatch(batchIdx); - SizeType32 const sequenceLength = sequenceLengthForBatch(batchIdx); - for (SizeType32 beamIdx = 0; beamIdx < beamWidth; ++beamIdx) - { - EXPECT_EQ(sequenceLengthsPtr[batchIdx * beamWidth + beamIdx], sequenceLength); - SizeType32 selectedCbaIdx = 0; - if (beamWidth == 1) - { - selectedCbaIdx = 1; - } - else - { - selectedCbaIdx = beamIdx == 0 ? 4 : (beamIdx == 1 ? 5 : 1); - } - - for (SizeType32 pos = 0; pos < sequenceLength; ++pos) - { - TokenIdType expected = promptToken(batchIdx, pos); - if (pos >= inputLength) - { - if (selectedCbaIdx >= beamWidth) - { - expected = unfinishedToken(batchIdx, selectedCbaIdx - beamWidth, pos); - } - else - { - expected = cbaToken(batchIdx, selectedCbaIdx, pos); - } - } - SizeType32 const dst = batchIdx * beamWidth * maxSeqLen + beamIdx * maxSeqLen + pos; - EXPECT_EQ(outputIdsPtr[dst], expected) - << "batchIdx=" << batchIdx << ", beamIdx=" << beamIdx << ", pos=" << pos; - } - } - } -} - -TEST(BeamHypothesesFinalizeTest, KeepsPromptPrefixFromUnfinishedIds) -{ - runFinalizeKeepsPromptPrefixFromUnfinishedIds(3); -} - -TEST(BeamHypothesesFinalizeTest, KeepsPromptPrefixFromUnfinishedIdsBeamWidthOne) -{ - runFinalizeKeepsPromptPrefixFromUnfinishedIds(1); -} - -/** - * @brief Fills a slice of a tensor with data from a source array. - * - * This function writes to `tensor` from source array `src` at index `idx. - * It optionally flattens the tensor before performing the insertion. - * For example tensor if we wanted to write 5 values in the 3rd row of [1,10,100] - * We will use (tensor, 2, 5, src, true, mBufferManager) where src is a buffer with at least 5 elems. - * - * @tparam T The type of elements in the source array. - * @param tensor A shared pointer to the tensor to be modified. Also need to be of type T. - * @param idx The index at which to start inserting data into the tensor. - * @param insertLen The number of elements to insert from the source array into the tensor. - * @param src An array containing the data to be inserted into the tensor. - * @param flattenFirst A boolean flag indicating whether to flatten the first dimension of the tensor before insertion. - * @param bufferManager A shared pointer to a BufferManager responsible for managing memory operations. - */ -template -void fillTensorAtIndex(ITensor::SharedPtr tensor, SizeType32 idx, std::vector src, bool flattenFirst, - std::shared_ptr bufferManager) -{ - SizeType32 insertLen = src.size(); - ITensor::SharedPtr target = ITensor::view(tensor); - if (flattenFirst) - { - target->squeeze(0); - } - - target = ITensor::slice(target, idx, 1); - target->squeeze(0); - target = ITensor::slice(target, 0, insertLen); - bufferManager->copy(src.data(), *target); -} - -} // anonymous namespace - -class TestGatherTree : public ::testing::Test -{ -public: - SizeType32 batchSize{1}; - SizeType32 beamWidth{5}; - SizeType32 maxSeqLen{20}; - - using TensorPtr = ITensor::SharedPtr; - - std::shared_ptr mStream{nullptr}; - std::shared_ptr mBufferManager{nullptr}; - - std::unique_ptr mDecodingState{nullptr}; - - SamplingConfig mSamplingConfig; - - TensorPtr mTargetOut{nullptr}; - - void SetUp() override - { - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); - } - - // create the empty buffers with the correct shapes and zero them - void createBuffers() - { - SizeType32 constexpr tensorParallelism{1}; - SizeType32 constexpr pipelineParallelism{1}; - SizeType32 constexpr contextParallelism{1}; - SizeType32 constexpr localRank{0}; - WorldConfig const worldConfig{tensorParallelism, pipelineParallelism, contextParallelism, localRank}; - - SizeType32 constexpr vocabSize{51200}; - SizeType32 constexpr nbAttentionLayers{2}; - SizeType32 constexpr nbRnnLayers{0}; - SizeType32 constexpr nbHeads{16}; - SizeType32 constexpr hiddenSize{1024}; - tensorrt_llm::DataType constexpr dtype{tensorrt_llm::DataType::kFLOAT}; - ModelConfig modelConfig{ - vocabSize, nbAttentionLayers + nbRnnLayers, nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, dtype}; - - mDecodingState = std::make_unique(); - mDecodingState->setup( - batchSize, beamWidth, maxSeqLen, 0, maxSeqLen, dtype, modelConfig, worldConfig, *mBufferManager); - - auto constexpr nvTokenIdType = TRTDataType::value; - auto const jointOutputIdsShape = ITensor::makeShape({batchSize, beamWidth, maxSeqLen}); - - mTargetOut = mBufferManager->gpu(jointOutputIdsShape, nvTokenIdType); - mBufferManager->setZero(*mTargetOut); - } - - // clang-format off - - // hardcode the input data for the output_len = 10 case - // this should not cause any beam swapping from the CBAs, just reorder the beams - void hardcodeBuffersLen10() - { - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - auto constexpr nvFloatType = TRTDataType::value; - - auto const decodingInput = mDecodingState->getJointDecodingInput(); - auto const decodingOutput = mDecodingState->getJointDecodingOutput(); - - std::vector len = {3, 3, 3, 3, 3}; - TensorPtr inputLengths{ITensor::slice(constPointerCast(decodingInput.lengths), 0, 1)}; - mBufferManager->copy(len.data(),*inputLengths); - - std::vector eid = {0}; - TensorPtr endIds{ITensor::slice(constPointerCast(decodingInput.endIds), 0, 1)}; - mBufferManager->copy(eid.data(),*endIds); - - std::vector> logProbs = - { - {-2.96689, -1.63675, -2.31329, -0.0377979, -2.2442, -1.57552, -0.310524, -0.696636, -2.41985}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -2.63339, -0.534199, -0.493615, -2.61479}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -2.2442, -1.57552, -0.310524, -3.11851, -1.01671}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -2.63339, -0.534199, 0, 0}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -2.2442, -1.57552, -0.310524, -0.696636, -3.62298} - }; - for (SizeType32 it = 0; it < logProbs.size(); it++){ - fillTensorAtIndex(decodingOutput.logProbs, it, logProbs[it], true, mBufferManager); - } - - std::vector> logProbsTiled = - { - {-2.70907, -2.96689, -3.27157, -3.37314, -3.50595}, - {-1.84733, -1.8942, -1.63675, -1.9567, -1.47513}, - {-0.305059, -0.765237, -2.31329, -2.37162, -2.48475}, - {-1.97517, -0.0377979, -2.0169, -2.42439, -2.27471}, - {-1.31451, -2.2442, -1.5831, -2.44732, -2.02409}, - {-1.57552, -2.63339, -2.11286, -2.57304, -3.85214}, - {-0.310524, -0.534199, -0.74379, -2.86232, -1.72914}, - {-0.696636, -0.493615, -0.237725, -3.07164, -3.11851}, - {-2.41985, -2.61479, -1.01671, -3.62298, -1.26586}, - {-0.844337, -0.922832, -0.427682, -0.419985, -1.85996} - }; - TensorPtr logProbsTiledView = ITensor::view(decodingOutput.logProbsTiled,ITensor::makeShape({maxSeqLen*batchSize, beamWidth})); - for (SizeType32 it = 0; it < logProbsTiled.size(); it++){ - auto logProbsSlice = ITensor::slice(logProbsTiledView, it+3,1); - mBufferManager->copy(logProbsTiled[it].data(),*logProbsSlice); - } - - std::vector outputLenghts = {13, 13, 13, 13, 13}; - mBufferManager->copy(outputLenghts.data(),*decodingOutput.lengths); - - std::vector cumLogProbs = {-15.0458, -15.4681, -15.8323, -15.8424, -16.0614}; - mBufferManager->copy(cumLogProbs.data(),*decodingOutput.cumLogProbs); - - std::vector> outputIdsCBA = - { - {1, 864, 304, 367, 263, 760, 310, 278, 3815, 29973}, - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973} - }; - for(SizeType32 it = 0; it < outputIdsCBA.size(); it++) - { - fillTensorAtIndex(decodingOutput.beamHypotheses.outputIdsCBA, it, outputIdsCBA[it], true, mBufferManager); - } - - std::vector> logProbsCBA = - { - {0, 0, 0, -2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -2.63339, -0.534199, -2.19674}, - {0, 0, 0, -2.96689, -1.63675, -2.31329, -0.0377979, -2.2442, -1.57552, -0.310524, -2.81382,} - }; - for(SizeType32 it = 0; it < logProbsCBA.size(); it++) - { - fillTensorAtIndex(decodingOutput.beamHypotheses.logProbsCBA, it, logProbsCBA[it], true, mBufferManager); - } - - std::vector sequenceLengthsCBA = {10, 10, 0, 0, 0, 0, 0, 0, 0, 0}; - mBufferManager->copy(sequenceLengthsCBA.data(), *decodingOutput.beamHypotheses.sequenceLengthsCBA); - - std::vector cumLogProbsCBA = {-13.6336, -13.8988, 0, 0, 0, 0, 0, 0, 0, 0}; - mBufferManager->copy(cumLogProbsCBA.data(), *decodingOutput.beamHypotheses.cumLogProbsCBA); - - std::vector normedScoresCBA = {-1.7042, -1.73735, 0, 0, 0, 0, 0, 0, 0, 0}; - mBufferManager->copy(normedScoresCBA.data(), *decodingOutput.beamHypotheses.normedScoresCBA); - - std::vector numBeamsCBA = {2}; - mBufferManager->copy(numBeamsCBA.data(), *decodingOutput.beamHypotheses.numBeamsCBA); - - std::vector minNormedScoresCBA = {-1.73735}; - mBufferManager->copy(minNormedScoresCBA.data(), *decodingOutput.beamHypotheses.minNormedScoresCBA); - - std::vector batchDones = {0}; - mBufferManager->copy(batchDones.data(), *decodingOutput.beamHypotheses.batchDones); - - std::vector finishReasons = {4, 4, 4, 4, 4}; - mBufferManager->copy(finishReasons.data(), *decodingOutput.finishReasons); - - std::vector> ids = - { - {1, 864, 304, 1073, 825, 1048, 278, 278, 3815, 29973, 13, 4806, 526}, - {1, 864, 304, 367, 920, 304, 310, 1749, 3815, 29973, 13, 4806, 526}, - {1, 864, 304, 679, 263, 760, 679, 263, 29973, 13, 310, 526, 502}, - {1, 864, 304, 1207, 901, 278, 1749, 445, 3889, 393, 591, 13443, 276}, - {1, 864, 304, 1074, 263, 29973, 1207, 263, 2446, 12623, 1334, 29915, 30010} - }; - for(SizeType32 it = 0; it < ids.size(); it++) - { - fillTensorAtIndex(decodingOutput.ids, it, ids[it], true, mBufferManager); - } - - std::vector> parentIds = - { - {0, 0, 0, 0, 0, 3, 0, 1, 1, 0, 0, 0, 0}, - {0, 0, 0, 0, 0, 1, 2, 1, 0, 1, 1, 1, 1}, - {0, 0, 0, 0, 1, 2, 1, 4, 3, 2, 4, 4, 3}, - {0, 0, 0, 0, 0, 0, 0, 1, 4, 1, 0, 0, 4}, - {0, 0, 0, 0, 3, 3, 1, 2, 0, 4, 0, 3, 0} - }; - for(SizeType32 it = 0; it < parentIds.size(); it++) - { - fillTensorAtIndex(decodingOutput.parentIds, it, parentIds[it], true, mBufferManager); - } - - std::vector> targetOutput = - { - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973, 13, 4806, 526}, - {1, 864, 304, 367, 263, 760, 310, 278, 3815, 29973, 13, 4806, 526}, - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973, 13, 13443, 502}, - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973, 591, 29915, 276}, - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973, 13, 4806, 30010} - }; - for(SizeType32 it = 0; it < targetOutput.size(); it++) - { - fillTensorAtIndex(mTargetOut, it, targetOutput[it], true, mBufferManager); - } - } - - // this case has the output_len = 8, and tests that the beams from the CBAs are correctly swapped. - void hardcodeBuffersLen8() - { - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - auto constexpr nvFloatType = TRTDataType::value; - - auto const decodingInput = mDecodingState->getJointDecodingInput(); - auto const decodingOutput = mDecodingState->getJointDecodingOutput(); - - std::vector len = {3, 3, 3, 3, 3}; - TensorPtr inputLengths{ITensor::slice(constPointerCast(decodingInput.lengths), 0, 1)}; - mBufferManager->copy(len.data(),*inputLengths); - - std::vector eid = {0}; - TensorPtr endIds{ITensor::slice(constPointerCast(decodingInput.endIds), 0, 1)}; - mBufferManager->copy(eid.data(),*endIds); - - std::vector >logProbs = - { - {-2.96689, -1.63675, -2.31329, -0.0377979, -2.2442, -1.57552, -0.310524}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -2.63339, -0.534199}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -2.44732, -2.11286, -0.74379}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -2.63339, -2.86232}, - {-2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -3.85214, -1.72914} - }; - for (SizeType32 it = 0; it < logProbs.size(); it++){ - fillTensorAtIndex(decodingOutput.logProbs, it, logProbs[it], true, mBufferManager); - } - - std::vector> logProbsTiled = - { - {-2.70907, -2.96689, -3.27157, -3.37314, -3.50595}, - {-1.84733, -1.8942, -1.63675, -1.9567, -1.47513}, - {-0.305059, -0.765237, -2.31329, -2.37162, -2.48475}, - {-1.97517, -0.0377979, -2.0169, -2.42439, -2.27471}, - {-1.31451, -2.2442, -1.5831, -2.44732, -2.02409}, - {-1.57552, -2.63339, -2.11286, -2.57304, -3.85214}, - {-0.310524, -0.534199, -0.74379, -2.86232, -1.72914}, - {-0.696636, -0.493615, -0.237725, -3.07164, -3.11851} - }; - TensorPtr logProbsTiledView = ITensor::view(decodingOutput.logProbsTiled,ITensor::makeShape({maxSeqLen*batchSize, beamWidth})); - for (SizeType32 it = 0; it < logProbsTiled.size(); it++){ - auto logProbsSlice = ITensor::slice(logProbsTiledView, it+3,1); - mBufferManager->copy(logProbsTiled[it].data(),*logProbsSlice); - } - std::vector outputLenghts = {11, 11, 11, 11, 11}; - mBufferManager->copy(outputLenghts.data(),*decodingOutput.lengths); - - std::vector cumLogProbs = {-11.7816, -11.9304, -14.0883, -14.1566, -14.2035}; - mBufferManager->copy(cumLogProbs.data(),*decodingOutput.cumLogProbs); - - std::vector> outputIdsCBA = - { - {1, 864, 304, 367, 263, 760, 310, 278, 3815, 29973}, - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973} - }; - for(SizeType32 it = 0; it < outputIdsCBA.size(); it++) - { - fillTensorAtIndex(decodingOutput.beamHypotheses.outputIdsCBA, it, outputIdsCBA[it], true, mBufferManager); - } - - std::vector> logProbsCBA = - { - {0, 0, 0, -2.96689, -1.63675, -2.31329, -0.0377979, -1.31451, -2.63339, -0.534199, -2.19674}, - {0, 0, 0, -2.96689, -1.63675, -2.31329, -0.0377979, -2.2442, -1.57552, -0.310524, -2.81382,} - }; - for(SizeType32 it = 0; it < logProbsCBA.size(); it++) - { - fillTensorAtIndex(decodingOutput.beamHypotheses.logProbsCBA, it, logProbsCBA[it], true, mBufferManager); - } - - std::vector sequenceLengthsCBA = {10, 10, 0, 0, 0, 0, 0, 0, 0, 0}; - mBufferManager->copy(sequenceLengthsCBA.data(), *decodingOutput.beamHypotheses.sequenceLengthsCBA); - - std::vector cumLogProbsCBA = {-13.6336, -13.8988, 0, 0, 0, 0, 0, 0, 0, 0}; - mBufferManager->copy(cumLogProbsCBA.data(), *decodingOutput.beamHypotheses.cumLogProbsCBA); - - std::vector normedScoresCBA = {-1.7042, -1.73735, 0, 0, 0, 0, 0, 0, 0, 0}; - mBufferManager->copy(normedScoresCBA.data(), *decodingOutput.beamHypotheses.normedScoresCBA); - - std::vector numBeamsCBA = {2}; - mBufferManager->copy(numBeamsCBA.data(), *decodingOutput.beamHypotheses.numBeamsCBA); - - std::vector minNormedScoresCBA = {-1.73735}; - mBufferManager->copy(minNormedScoresCBA.data(), *decodingOutput.beamHypotheses.minNormedScoresCBA); - - std::vector batchDones = {0}; - mBufferManager->copy(batchDones.data(), *decodingOutput.beamHypotheses.batchDones); - - std::vector finishReasons = {4, 4, 4, 4, 4}; - mBufferManager->copy(finishReasons.data(), *decodingOutput.finishReasons); - - std::vector> ids = - { - {1, 864, 304, 1073, 825, 1048, 278, 278, 3815, 29973, 13}, - {1, 864, 304, 367, 920, 304, 310, 1749, 3815, 29973, 13}, - {1, 864, 304, 679, 263, 760, 679, 263, 29973, 13, 310}, - {1, 864, 304, 1207, 901, 278, 1749, 445, 3889, 393, 591}, - {1, 864, 304, 1074, 263, 29973, 1207, 263, 2446, 12623, 1334} - }; - for(SizeType32 it = 0; it < ids.size(); it++) - { - fillTensorAtIndex(decodingOutput.ids, it, ids[it], true, mBufferManager); - } - - std::vector> parentIds = - { - {0, 0, 0, 0, 0, 3, 0, 1, 1, 0, 0}, - {0, 0, 0, 0, 0, 1, 2, 1, 0, 1, 1}, - {0, 0, 0, 0, 1, 2, 1, 4, 3, 2, 4}, - {0, 0, 0, 0, 0, 0, 0, 1, 4, 1, 0}, - {0, 0, 0, 0, 3, 3, 1, 2, 0, 4, 0} - }; - for(SizeType32 it = 0; it < parentIds.size(); it++) - { - fillTensorAtIndex(decodingOutput.parentIds, it, parentIds[it], true, mBufferManager); - } - - std::vector> targetOutput = - { - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973, 13}, - {1, 864, 304, 367, 263, 760, 310, 278, 3815, 29973, 13}, - {1, 864, 304, 367, 263, 760, 310, 278, 3815, 29973, 0}, - {1, 864, 304, 367, 263, 760, 310, 1749, 3815, 29973, 0}, - {1, 864, 304, 367, 263, 760, 310, 278, 2446, 12623, 310} - }; - for(SizeType32 it = 0; it < targetOutput.size(); it++) - { - fillTensorAtIndex(mTargetOut, it, targetOutput[it], true, mBufferManager); - } - } - - // clang-format on - - bool checkResult() - { - auto const reference = this->mBufferManager->copyFrom(*mTargetOut, tensorrt_llm::runtime::MemoryType::kCPU); - auto referencePtr = bufferCast(*reference); - - auto const real = this->mBufferManager->copyFrom( - *mDecodingState->getGatheredIds(), tensorrt_llm::runtime::MemoryType::kCPU); - auto realPtr = bufferCast(*real); - - bool allEqual = true; - for (SizeType32 iAssert = 0; iAssert < batchSize * beamWidth * maxSeqLen; iAssert++) - { - if (referencePtr[iAssert] != realPtr[iAssert]) - { - TLLM_LOG_ERROR("Mismatch input value. Position of inputs: %d, expected value: %d, output value: %d", - iAssert, referencePtr[iAssert], realPtr[iAssert]); - allEqual = false; - } - } - return allEqual; - } -}; - -TEST_F(TestGatherTree, GatherTreeNoSwap) -{ - createBuffers(); - hardcodeBuffersLen10(); - cudaDeviceSynchronize(); - kernels::gatherTree( - mDecodingState->getJointDecodingOutput(), mDecodingState->getJointDecodingInput(), mSamplingConfig, *mStream); - cudaDeviceSynchronize(); - - EXPECT_TRUE(checkResult()); -} - -TEST_F(TestGatherTree, GatherTreeWithSwap) -{ - createBuffers(); - hardcodeBuffersLen8(); - cudaDeviceSynchronize(); - kernels::gatherTree( - mDecodingState->getJointDecodingOutput(), mDecodingState->getJointDecodingInput(), mSamplingConfig, *mStream); - cudaDeviceSynchronize(); - - EXPECT_TRUE(checkResult()); -} - -// Test that generation logits are correctly reordered after gatherTree finalization. -// Uses the same hardcoded beam search data as GatherTreeNoSwap, creates sentinel logits -// where logits[slot][g][v] = slot (so we can verify which slot each beam's logits came from), -// then runs the reorder algorithm and checks the result. -TEST_F(TestGatherTree, GenerationLogitsReorder) -{ - // Compile-time proof that the fixture data causes beam reordering. - // hardcodeBuffersLen10: parentIds[slot=1][t=3] = 0 (row 1, col 3 of the parentIds table). - // Since 0 != 1, any beam trace through slot 1 at t=4 steps to slot 0 at t=3 — a concrete swap. - static constexpr SizeType32 kFixtureParentIds_slot1_t3 = 0; - static_assert(kFixtureParentIds_slot1_t3 != 1, - "parentIds[slot=1][t=3] must differ from slot 1 to guarantee a beam swap in this test"); - - createBuffers(); - hardcodeBuffersLen10(); - cudaDeviceSynchronize(); - kernels::gatherTree( - mDecodingState->getJointDecodingOutput(), mDecodingState->getJointDecodingInput(), mSamplingConfig, *mStream); - cudaDeviceSynchronize(); - - // Verify gatherTree worked first - ASSERT_TRUE(checkResult()); - - // Now test the generation logits reordering algorithm. - // Copy ids, parentIds, gatheredIds, and seqLengths to host. - auto idsHost = mBufferManager->copyFrom(*mDecodingState->getIds(0), MemoryType::kCPU); - auto parentIdsHost = mBufferManager->copyFrom(*ITensor::at(mDecodingState->getParentIds(), {0}), MemoryType::kCPU); - auto gatheredIdsHost = mBufferManager->copyFrom(*mDecodingState->getGatheredIds(0), MemoryType::kCPU); - auto seqLengthsHost = mBufferManager->copyFrom(*mDecodingState->getSequenceLengths(0), MemoryType::kCPU); - - auto const* idsData = bufferCast(*idsHost); - auto const* parentIdsData = bufferCast(*parentIdsHost); - auto const* gatheredIdsData = bufferCast(*gatheredIdsHost); - auto const* seqLengthsData = bufferCast(*seqLengthsHost); - - SizeType32 constexpr promptLen = 3; // matches inputLengths in hardcodeBuffersLen10 - SizeType32 constexpr vocabSizePadded = 32000; // LLaMA vocab size (matches fixture token IDs) - SizeType32 const maxNewTokens = maxSeqLen - promptLen; - - // Populate sentinel logits: for each (pre-reassignment slot, gen step), place a 1.0f - // at the selected token position. All other entries stay 0. After correct reordering, - // logits[beam][g][gatheredToken] should be positive (the sentinel landed at the right place). - std::vector logits(static_cast(beamWidth) * maxNewTokens * vocabSizePadded, 0.0f); - for (SizeType32 postSlot = 0; postSlot < beamWidth; ++postSlot) - { - auto const genLen = seqLengthsData[postSlot] - promptLen; - for (SizeType32 g = 0; g < genLen; ++g) - { - SizeType32 const t = promptLen + g; - SizeType32 const preSlot = parentIdsData[postSlot * maxSeqLen + t]; - TokenIdType const token = idsData[postSlot * maxSeqLen + t]; - logits[static_cast(preSlot * maxNewTokens + g) * vocabSizePadded + token] = 1.0f; - } - } - - // Build slot trace and reorder (same algorithm as reorderGenerationLogitsForBeamSearch) - std::vector> slotTrace(beamWidth, std::vector(maxNewTokens, 0)); - - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto const seqLen = seqLengthsData[beam]; - auto const genLen = seqLen - promptLen; - if (genLen <= 0) - { - continue; - } - - // Find starting slot by matching backtracked sequence - SizeType32 startSlot = -1; - for (SizeType32 s = 0; s < beamWidth; ++s) - { - SizeType32 slot = s; - bool matches = true; - for (SizeType32 t = seqLen - 1; t >= promptLen; --t) - { - if (idsData[slot * maxSeqLen + t] != gatheredIdsData[beam * maxSeqLen + t]) - { - matches = false; - break; - } - if (t > promptLen) - { - slot = parentIdsData[slot * maxSeqLen + t]; - } - } - if (matches) - { - startSlot = s; - break; - } - } - ASSERT_GE(startSlot, 0) << "Could not find starting slot for beam " << beam; - - // Build pre-reassignment slot trace in a single pass - SizeType32 slot = startSlot; - for (SizeType32 t = seqLen - 1; t >= promptLen; --t) - { - slot = parentIdsData[slot * maxSeqLen + t]; - slotTrace[beam][t - promptLen] = slot; - } - } - - // Reorder logits using a temp buffer (same approach as the production code) - auto const stepSize = static_cast(vocabSizePadded) * sizeof(float); - std::vector temp(beamWidth * vocabSizePadded); - auto* logitsPtr = reinterpret_cast(logits.data()); - auto* tempPtr = reinterpret_cast(temp.data()); - - std::vector genLens(beamWidth); - SizeType32 maxGenLen = 0; - for (SizeType32 b = 0; b < beamWidth; ++b) - { - genLens[b] = std::max(SizeType32{0}, seqLengthsData[b] - promptLen); - maxGenLen = std::max(maxGenLen, genLens[b]); - } - - for (SizeType32 g = 0; g < maxGenLen; ++g) - { - bool stepNeedsReorder = false; - for (SizeType32 b = 0; b < beamWidth; ++b) - { - if (g < genLens[b] && slotTrace[b][g] != b) - { - stepNeedsReorder = true; - break; - } - } - if (!stepNeedsReorder) - { - continue; - } - - for (SizeType32 b = 0; b < beamWidth; ++b) - { - auto const offset = (static_cast(b) * maxNewTokens + g) * stepSize; - std::memcpy(tempPtr + static_cast(b) * stepSize, logitsPtr + offset, stepSize); - } - for (SizeType32 b = 0; b < beamWidth; ++b) - { - if (g >= genLens[b]) - { - continue; - } - auto const dstOffset = (static_cast(b) * maxNewTokens + g) * stepSize; - auto const srcSlot = slotTrace[b][g]; - std::memcpy(logitsPtr + dstOffset, tempPtr + static_cast(srcSlot) * stepSize, stepSize); - } - } - - // Cross-check: after reorder, logits[beam][g][gatheredToken] should equal the 1.0f sentinel - // placed there during population. This is non-tautological: gatheredIdsData is produced - // independently by gatherTree tracing through parentIds, while the logits were reordered via - // the slot trace. A wrong slot's logits would have 0.0f at this position. - bool allCorrect = true; - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto const genLen = genLens[beam]; - for (SizeType32 g = 0; g < genLen; ++g) - { - TokenIdType const gatheredToken = gatheredIdsData[beam * maxSeqLen + (promptLen + g)]; - float const logitAtGatheredToken - = logits[static_cast(beam * maxNewTokens + g) * vocabSizePadded + gatheredToken]; - if (logitAtGatheredToken != 1.0f) - { - TLLM_LOG_ERROR("Beam %d, step %d: logit at gathered token %d is %.1f, expected 1.0", beam, g, - gatheredToken, logitAtGatheredToken); - allCorrect = false; - } - } - } - EXPECT_TRUE(allCorrect); -} - -namespace -{ - -enum AcceptKernelMode -{ - BY_IDS, - BY_LOGITS, - BY_IDS_WITH_PATH -}; - -struct DecodingKernelTestParam -{ - SizeType32 mBatchSize{128}; - SizeType32 mMaxBatchSize{2 * mBatchSize}; - SizeType32 mBeamWidth{1}; - SizeType32 mMaxSeqLen{16}; - SizeType32 mVocabSize{32}; - SizeType32 mMaxDraftTokens{8}; - SizeType32 mMaxNumHeads{0}; - SizeType32 mMaxDraftSeqPerStep{1}; - AcceptKernelMode mAcceptMode{AcceptKernelMode::BY_IDS}; - - DecodingKernelTestParam& setBatchSize(SizeType32 bs) - { - mBatchSize = bs; - mMaxBatchSize = 2 * mBatchSize; - return *this; - } - - DecodingKernelTestParam& setVocabSize(SizeType32 vs) - { - mVocabSize = vs; - return *this; - } - - DecodingKernelTestParam& setMaxSeqLen(SizeType32 msl) - { - mMaxSeqLen = msl; - return *this; - } - - DecodingKernelTestParam& setMaxDraftTokens(SizeType32 dt) - { - mMaxDraftTokens = dt; - return *this; - } - - DecodingKernelTestParam& setMaxNumHeads(SizeType32 mnh) - { - mMaxNumHeads = mnh; - return *this; - } - - DecodingKernelTestParam& setMaxDraftSeqPerStep(SizeType32 tps) - { - mMaxDraftSeqPerStep = tps; - return *this; - } - - DecodingKernelTestParam& setAcceptMode(AcceptKernelMode const& mode) - { - mAcceptMode = mode; - return *this; - } -}; - -template -class DecodingKernelsTest : public testing::Test -{ -public: - using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; - - void SetUp() override - { - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); - } - - void TearDown() override {} - - void createBuffers() - { - auto const dataType = TRTDataType::value; - auto const ptrType = TRTDataType::value; - - mDraftTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqlen}), tensorrt_llm::DataType::kINT32); - mTargetTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTargetSeqlen}), tensorrt_llm::DataType::kINT32); - mOutputTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mNumsDraftTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep}), tensorrt_llm::DataType::kINT32); - mSequenceLengths - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mAcceptedLengths - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mContextLengths - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mFinishedSteps = mBufferManager->pinnedPool(ITensor::makeShape({mMaxDraftTokens + 1, mMaxBatchSize}), - TRTDataType::value); - mFinishedFinal = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize}), TRTDataType::value); - mFinishedSum = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mPaths = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep, mMaxDraftTokens}), tensorrt_llm::DataType::kINT32); - mEndIds = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mBatchSlots = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - auto batchSlotsRange = BufferRange(*mBatchSlots); - std::iota(batchSlotsRange.begin(), batchSlotsRange.end(), 0); - - mCurandStates = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); - - mAcceptedLen.resize(mMaxBatchSize); - mOutputLen.resize(mMaxBatchSize); - mAcceptedFinished.resize(mMaxBatchSize, tk::FinishedState::empty()); - - // Buffers only for Logits comparison - if (mAcceptMode == AcceptKernelMode::BY_LOGITS) - { - mDraftLogits = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTotalDraftTokens, mVocabSize}), dataType); - mTargetLogits = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTotalDraftTokens, mVocabSize}), dataType); - mTargetLogitsPtrs = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), ptrType); - mRefTargetLogits = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTotalDraftTokens, mVocabSize}), dataType); - - mDraftProbs = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTotalDraftTokens, mVocabSize}), dataType); - mTargetProbs = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTotalDraftTokens, mVocabSize}), dataType); - } - - if (mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH) - { - mMedusaLogitsPtrs = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep, mMaxNumHeads}), ptrType); - mMedusaInputLogitsPtrs - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize, mMaxNumHeads}), ptrType); - mTokensPerStep - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mBestPaths - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - } - } - - void initData(SizeType32 seed) - { - std::mt19937 generator(seed); - std::uniform_int_distribution contextLenDistr(0, std::max(mMaxSeqLen - mMaxTotalDraftTokens, 0)); - std::uniform_int_distribution numTotalDraftTokensDistr(1, mMaxTotalDraftTokens); - std::uniform_int_distribution numDraftTokensDistr(0, mMaxDraftTokens); - std::uniform_int_distribution vocabDistr(1, mVocabSize - 1); - std::uniform_real_distribution acceptTokenDistr(0.f, 1.f); - - trk::invokeFill(*mPaths, int32_t{-1}, *mStream); - trk::invokeFill(*mFinishedFinal, tk::FinishedState::UnderlyingType{0}, *mStream); - - auto sequenceLengthsPtr = BufferRange(*mSequenceLengths); - auto contextLengthsPtr = BufferRange(*mContextLengths); - auto numsDraftTokensPtr = BufferRange(*mNumsDraftTokens); - auto draftTokensPtr = BufferRange(*mDraftTokens); - auto targetTokensPtr = BufferRange(*mTargetTokens); - auto finishedStepsPtr - = reinterpret_cast(bufferCast(*mFinishedSteps)); - auto pathsPtr = BufferRange(*mPaths); - auto endIdsPtr = BufferRange(*mEndIds); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - - tk::invokeCurandInitialize(reinterpret_cast(bufferCast(*mCurandStates)), batchSlotsPtr, - mMaxBatchSize, seed, this->mStream->get()); - - auto generateAvoidingValues = [&vocabDistr, &generator](std::uniform_int_distribution& distr, - std::unordered_set const& tokensToAvoid, SizeType32 maxTries = -1, - SizeType32 defaultValue = -1) - { - // Avoid generating endId. - auto token = distr(generator); - SizeType32 tries = 0; - while (tokensToAvoid.count(token) != 0 && ((maxTries >= 0 && tries < maxTries) || maxTries < 0)) - { - token = distr(generator); - tries++; - } - if (tries == maxTries) - { - token = defaultValue; - } - return token; - }; - - // Init batch slots - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - batchSlotsPtr[bi] = 2 * bi; - } - - // Init end ids - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - endIdsPtr[bi] = generateAvoidingValues(vocabDistr, {mPadId}); - TLLM_LOG_DEBUG("bi %d endIdsPtr[bi] %d", bi, endIdsPtr[bi]); - - // Randomly init context len for target and draft - contextLengthsPtr[bi] = contextLenDistr(generator); - } - - std::fill(draftTokensPtr.begin(), draftTokensPtr.begin() + mMaxBatchSize * mMaxDraftSeqlen, mPadId); - std::fill(targetTokensPtr.begin(), targetTokensPtr.begin() + mMaxBatchSize * mMaxTargetSeqlen, mPadId); - std::fill(pathsPtr.begin(), pathsPtr.begin() + mMaxBatchSize * mMaxDraftSeqPerStep * mMaxDraftTokens, -1); - - // Generate paths - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - auto const numTotalDraftTokens = std::min(mMaxDraftTokens, numTotalDraftTokensDistr(generator)); - std::uniform_int_distribution pathIdDistr(0, numTotalDraftTokens); - for (SizeType32 pi = 0; pi < mMaxDraftSeqPerStep; ++pi) - { - std::unordered_set pathIds; - auto const numDraftTokensAtStep = numDraftTokensDistr(generator); - numsDraftTokensPtr[bi * mMaxDraftSeqPerStep + pi] = numDraftTokensAtStep; - - for (SizeType32 ti = 0; ti < numDraftTokensAtStep; ++ti) - { - auto const pathIdx = tc::flat_index3(bi, pi, ti, mMaxDraftSeqPerStep, mMaxDraftTokens); - // Single linear path for BY_IDS and BY_LOGITS modes - auto const pathId = mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH && ti != 0 - ? generateAvoidingValues(pathIdDistr, pathIds, mMaxDraftTokens * 5, -1) - : ti; - pathsPtr[pathIdx] = pathId; - pathIds.insert(pathId); - } - if (bi == 2) - { - TLLM_LOG_DEBUG("bi %d pi %d numsDraftTokensPtr[bi] %d", bi, pi, numDraftTokensAtStep); - } - } - } - - for (SizeType32 ti = 0; ti < mMaxDraftSeqPerStep; ++ti) - { - std::vector targetPredictedLen(mMaxBatchSize); - std::vector targetAcceptedLen(mMaxBatchSize); - - // Init number of draft tokens - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - // It can be shorter than num of draft tokens due to the EOS generation - std::uniform_int_distribution realDraftTokensDistr( - 0, numsDraftTokensPtr[bi * mMaxDraftSeqPerStep + ti]); - targetPredictedLen[bi] = realDraftTokensDistr(generator); - // Accept ~ half of the tokens on avergae - std::poisson_distribution targetAcceptedDistr(targetPredictedLen[bi] / 2); - targetAcceptedLen[bi] = std::min(targetAcceptedDistr(generator), targetPredictedLen[bi]); - if (bi == 2) - { - TLLM_LOG_DEBUG("bi %d ti %d targetPredictedLen[bi] %d targetAcceptedLen[bi] %d", bi, ti, - targetPredictedLen[bi], targetAcceptedLen[bi]); - } - } - - // Fill draft tokens - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - for (SizeType32 si = 0; si < numsDraftTokensPtr[bi * mMaxDraftSeqPerStep + ti]; ++si) - { - auto const pathIdx = tc::flat_index3(bi, ti, si, mMaxDraftSeqPerStep, mMaxDraftTokens); - if (pathsPtr[pathIdx] == -1) - { - continue; - } - auto const draftTokenIdx = bi * mMaxDraftSeqlen + pathsPtr[pathIdx]; - // Avoid generating endId. We'll insert in manually later if needed. - draftTokensPtr[draftTokenIdx] = generateAvoidingValues(vocabDistr, {mPadId, endIdsPtr[bi]}); - if (bi == 2) - { - TLLM_LOG_DEBUG("bi %d ti %d si %d pathId %d draftToken %d", bi, ti, si, pathsPtr[pathIdx], - draftTokensPtr[draftTokenIdx]); - } - } - } - - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - sequenceLengthsPtr[bi] = contextLengthsPtr[bi] + targetPredictedLen[bi]; - - // Initialize finished states - for (int di = 0; di < numsDraftTokensPtr[bi * mMaxDraftSeqPerStep + ti]; ++di) - { - finishedStepsPtr[di * mMaxBatchSize + bi] - = (di < targetPredictedLen[bi]) ? tk::FinishedState::empty() : tk::FinishedState::finished(); - } - - // Init helper vectors - mAcceptedLen[bi] = contextLengthsPtr[bi] + std::max(targetAcceptedLen[bi], 0); - mOutputLen[bi] = std::min(sequenceLengthsPtr[bi], std::min(mAcceptedLen[bi] + 1, mMaxSeqLen)); - mAcceptedFinished[bi] = finishedStepsPtr[std::max(targetAcceptedLen[bi], 0) * mMaxBatchSize + bi]; - if (bi == 2) - { - TLLM_LOG_DEBUG( - "bi %d ti %d contextLengthsPtr[bi] %d sequenceLengthsPtr[bi] %d mAcceptedLen[bi] %d " - "mOutputLen[bi] " - "%d", - bi, ti, contextLengthsPtr[bi], sequenceLengthsPtr[bi], mAcceptedLen[bi], mOutputLen[bi]); - } - } - - // Fill token arrays - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - // Draft: [d0, d1, d2, ... for numsDraftTokensPtr[bi] ... , dK, - // padId, padId, .. to mMaxDraftSeqlen] - // Target: [padId, padId, ... for contextLengthsPtr[bi] ... padId, - // d0, d1, d2, ... for targetAcceptedLen[bi], - // ti (!= di), ti+1 (!= di+1), ... for (targetPredictedLen[bi] - targetAcceptedLen[bi]), - // EOS, EOS, EOS, ... for (numsDraftTokensPtr[bi] - targetPredictedLen[bi]) - // padId, padId, .. to mMaxSeqLen] - auto numDraftTokens = numsDraftTokensPtr[bi * mMaxDraftSeqPerStep + ti]; - for (SizeType32 si = 0; si < numDraftTokens; ++si) - { - auto const curPathIdx = tc::flat_index3(bi, ti, si, mMaxDraftSeqPerStep, mMaxDraftTokens); - auto const nextPathIdx = si + 1 < numDraftTokens - ? tc::flat_index3(bi, ti, si + 1, mMaxDraftSeqPerStep, mMaxDraftTokens) - : -1; - auto const curPathId = pathsPtr[curPathIdx]; - auto nextPathId = curPathId; - if (mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH) - { - nextPathId = nextPathIdx > -1 ? pathsPtr[nextPathIdx] : -1; - } - - if (curPathId == -1) - { - continue; - } - auto const contextLen - = mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH ? 0 : contextLengthsPtr[bi]; - auto const draftTokenIdx = bi * mMaxDraftSeqlen + nextPathId; - auto const targetTokenIdx = bi * mMaxTargetSeqlen + contextLen + curPathId; - auto targetToken = mPadId; - if (0 <= si && si < targetAcceptedLen[bi] && nextPathId != -1) - { - // Use draft token up to the accepted len - targetToken = draftTokensPtr[draftTokenIdx]; - } - else if (0 <= si && si < targetPredictedLen[bi]) - { - // Do not use draft token token up to the generated len - std::unordered_set avoidValues = {mPadId, endIdsPtr[bi]}; - if (nextPathId != -1) - { - avoidValues.insert(draftTokensPtr[draftTokenIdx]); - } - targetToken = generateAvoidingValues(vocabDistr, avoidValues); - } - else if (targetPredictedLen[bi] <= si && si < numsDraftTokensPtr[bi]) - { - // Fill with EOS from generated len to the draft len - targetToken = endIdsPtr[bi]; - } - targetTokensPtr[targetTokenIdx] = targetToken; - if (bi == 2) - { - TLLM_LOG_DEBUG( - "bi %d ti %d si %d pathId %d targetToken %d", bi, ti, si, curPathId, targetToken); - } - } - } - } - - if (mAcceptMode == AcceptKernelMode::BY_LOGITS) - { - initDataAndReferenceAcceptByLogits(); - } - - if (mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH) - { - initDataAndReferenceAcceptByIdsWithPaths(); - } - mSequenceLengthsCopy = mBufferManager->copyFrom(*mSequenceLengths, MemoryType::kCPU); - } - - void initDataAndReferenceAcceptByIdsWithPaths() - { - auto const dataType = TRTDataType::value; - auto const ptrType = TRTDataType::value; - - auto pathsPtr = BufferRange(*mPaths); - auto endIdsPtr = BufferRange(*mEndIds); - auto contextLengthsPtr = BufferRange(*mContextLengths); - auto draftTokensPtr = BufferRange(*mDraftTokens); - auto targetTokensPtr = BufferRange(*mTargetTokens); - auto medusaInputLogitsPtr = BufferRange(*mMedusaInputLogitsPtrs); - - trk::invokeFill(*mMedusaLogitsPtrs, int64_t{0}, *mStream); - trk::invokeFill(*mTokensPerStep, int32_t{mMaxTotalDraftTokens}, *mStream); - trk::invokeFill(*mBestPaths, int32_t{-1}, *mStream); - - mAcceptedLen.resize(mMaxBatchSize); - mAcceptedPathIdx.resize(mMaxBatchSize); - mRefAcceptedTokens.resize(mMaxBatchSize); - mFinishedByIdsPaths.resize(mMaxBatchSize); - mLastTargetIdx.resize(mMaxBatchSize); - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - SizeType32 maxAcceptedLen = -1; - SizeType32 maxAcceptedPath = -1; - SizeType32 maxNextTargetTokenIdx = -1; - bool maxFinished = false; - std::vector maxAcceptedTokens; - for (SizeType32 ti = 0; ti < mMaxDraftSeqPerStep; ++ti) - { - std::vector acceptedTokens; - SizeType32 curAcceptedLen = mMaxDraftTokens; - SizeType32 curAcceptedPath = ti; - bool curFinished = false; - - auto const pathIdx = tc::flat_index3(bi, ti, 0, mMaxDraftSeqPerStep, mMaxDraftTokens); - auto const pathId = pathsPtr[pathIdx]; - if (pathId == -1) - { - continue; - } - auto targetTokenIdx = bi * mMaxTargetSeqlen + pathId; - auto targetToken = targetTokensPtr[targetTokenIdx]; - auto curNextTargetTokenIdx = pathId; - for (SizeType32 di = 1; di < mMaxDraftTokens; ++di) - { - auto const pathIdx = tc::flat_index3(bi, ti, di, mMaxDraftSeqPerStep, mMaxDraftTokens); - auto const pathId = pathsPtr[pathIdx]; - if (pathId == -1) - { - curAcceptedLen = di; - curAcceptedPath = ti; - curFinished = false; - acceptedTokens.push_back(targetToken); - break; - } - auto const draftTokenIdx = bi * mMaxDraftSeqlen + pathId - 1; - auto const targetTokenIdx = bi * mMaxTargetSeqlen + pathId; - auto const draftToken = draftTokensPtr[draftTokenIdx]; - bool const hasEnd = targetToken == endIdsPtr[bi]; - if (!hasEnd) - { - acceptedTokens.push_back(targetToken); - } - if (draftToken != targetToken || hasEnd) - { - auto const curLen = hasEnd ? di - 1 : di; - curAcceptedLen = curLen; - curAcceptedPath = ti; - curFinished = hasEnd; - break; - } - targetToken = targetTokensPtr[targetTokenIdx]; - curNextTargetTokenIdx = pathId; - } - if (curAcceptedLen == mMaxDraftTokens) - { - acceptedTokens.push_back(targetToken); - } - if (curAcceptedLen > maxAcceptedLen) - { - maxAcceptedLen = curAcceptedLen; - maxAcceptedPath = curAcceptedPath; - maxAcceptedTokens = acceptedTokens; - maxFinished = curFinished; - maxNextTargetTokenIdx = curNextTargetTokenIdx; - } - } - mAcceptedLen[bi] = maxAcceptedLen; - mAcceptedPathIdx[bi] = maxAcceptedPath; - mRefAcceptedTokens[bi] = maxAcceptedTokens; - mFinishedByIdsPaths[bi] = maxFinished; - mLastTargetIdx[bi] = maxNextTargetTokenIdx; - for (SizeType32 hi = 0; hi < mMaxNumHeads; ++hi) - { - medusaInputLogitsPtr[bi * mMaxNumHeads + hi] = static_cast(nullptr) - + tc::flat_index4(hi, bi, 0, 0, mMaxBatchSize, mMaxDraftSeqPerStep, mVocabSize); - } - if (bi == 2) - { - TLLM_LOG_DEBUG("bi %d maxAcceptedLen %d maxAcceptedPath %d maxNextTargetTokenIdx %d", bi, - maxAcceptedLen, maxAcceptedPath, maxNextTargetTokenIdx); - std::ostringstream ss; - for (auto& tk : maxAcceptedTokens) - { - ss << tk << " "; - } - TLLM_LOG_DEBUG(ss.str().c_str()); - } - } - } - - void initDataAndReferenceAcceptByLogits() - { - auto contextLengthsPtr = BufferRange(*mContextLengths); - auto numsDraftTokensPtr = BufferRange(*mNumsDraftTokens); - auto draftTokensPtr = BufferRange(*mDraftTokens); - auto targetTokensPtr = BufferRange(*mTargetTokens); - - auto draftProbsPtr = BufferRange(*mDraftProbs); - auto targetProbsPtr = BufferRange(*mTargetProbs); - - auto draftLogitsPtr = BufferRange(*mDraftLogits); - auto targetLogitsPtr = BufferRange(*mTargetLogits); - auto targetLogitsPtrsPtr = BufferRange(*mTargetLogitsPtrs); - auto refTargetLogitsPtr = BufferRange(*mRefTargetLogits); - auto batchSlotsPtr = BufferRange(*mBatchSlots); - - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - // Init draft and target logits and probabilities - for (SizeType32 si = 0; si < numsDraftTokensPtr[bi]; ++si) - { - std::vector peakDraftProb(mVocabSize, 0.f); - std::vector peakTargetProb(mVocabSize, 0.f); - - auto const targetToken = targetTokensPtr[bi * mMaxSeqLen + contextLengthsPtr[bi] + si] % mVocabSize; - auto const draftToken = draftTokensPtr[bi * mMaxDraftTokens + si] % mVocabSize; - - peakDraftProb[draftToken] = 1.f; - peakTargetProb[targetToken] = 1.f; - - auto const logitsOffset = bi * mMaxDraftTokens * mVocabSize + si * mVocabSize; - // Emulate some distribution around target token - applyGaussianFilter( - draftProbsPtr.begin() + logitsOffset, peakDraftProb.data(), peakDraftProb.size(), 1.0f); - applyGaussianFilter( - targetProbsPtr.begin() + logitsOffset, peakTargetProb.data(), peakTargetProb.size(), 1.0f); - - // Probabilities to logits - probsToLogits(draftProbsPtr.begin() + logitsOffset, draftLogitsPtr.begin() + logitsOffset, mVocabSize); - probsToLogits( - targetProbsPtr.begin() + logitsOffset, targetLogitsPtr.begin() + logitsOffset, mVocabSize); - - // Do softmax conversion back to emulate kernels accuracy - softmax(draftLogitsPtr.begin() + logitsOffset, draftProbsPtr.begin() + logitsOffset, mVocabSize); - softmax(targetLogitsPtr.begin() + logitsOffset, targetProbsPtr.begin() + logitsOffset, mVocabSize); - } - } - - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - for (SizeType32 si = 0; si < mMaxDraftTokens; ++si) - { - auto const logitsOffset = bi * mMaxDraftTokens * mVocabSize + si * mVocabSize; - auto const outputLen = mOutputLen[bi] - contextLengthsPtr[bi]; - auto const acceptedLen = mAcceptedLen[bi] - contextLengthsPtr[bi]; - if (si < acceptedLen) - { - auto logitsStart = targetLogitsPtr.begin() + logitsOffset; - std::copy(logitsStart, logitsStart + mVocabSize, refTargetLogitsPtr.begin() + logitsOffset); - } - else if (si == acceptedLen) - { - // When token is not accepted, correct probabilities and compute updated logits - float sumProb = 1e-6f; - for (SizeType32 vi = 0; vi < mVocabSize; ++vi) - { - auto const correctedProb = std::max( - static_cast(targetProbsPtr[logitsOffset + vi] - draftProbsPtr[logitsOffset + vi]), - 0.f); - sumProb += correctedProb; - } - for (SizeType32 vi = 0; vi < mVocabSize; ++vi) - { - auto prob = std::max(static_cast( - targetProbsPtr[logitsOffset + vi] - draftProbsPtr[logitsOffset + vi]), - 0.f) - / sumProb; - if (prob < 1e-8) - { - prob = 0.f; - } - refTargetLogitsPtr[logitsOffset + vi] = std::log(prob / (1.f - prob)); - } - } - } - } - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - targetLogitsPtrsPtr[bi] = targetLogitsPtr.begin() + batchSlotsPtr[bi] * mMaxDraftTokens * mVocabSize; - } - } - - void callAcceptByIds() - { - // tksp::invokeAcceptDraftTokensByIds(bufferCast(*mDraftTokens), - // bufferCast(*mTargetTokens), bufferCast(*mContextLengths), - // bufferCast(*mNumsDraftTokens), bufferCast(*mSequenceLengths), - // reinterpret_cast(bufferCast(*mFinishedSteps)), - // reinterpret_cast(bufferCast(*mFinishedFinal)), - // bufferCast(*mFinishedSum), bufferCast(*mBatchSlots), mBatchSize, mMaxBatchSize, - // mBeamWidth, mMaxSeqLen, mMaxDraftTokens, mStream->get()); - } - - void callAcceptByLogits() - { - // tksp::acceptDraftTokensByLogits(bufferCast(*mDraftLogits), - // reinterpret_cast(bufferCast(*mTargetLogitsPtrs)), bufferCast(*mDraftProbs), - // bufferCast(*mTargetProbs), bufferCast(*mNumsDraftTokens), - // reinterpret_cast(bufferCast(*mFinishedSteps)), - // reinterpret_cast(bufferCast(*mCurandStates)), - // bufferCast(*mBatchSlots), mBatchSize, mMaxBatchSize, mBeamWidth, mVocabSize, mVocabSize, - // mMaxDraftTokens, false, 0.9f, mStream->get()); - } - - void callAcceptByIdsWithPaths() - { - tksp::AcceptDraftTokensByIdsWithPathsParams params; - - params.outputIds = bufferCast(*mOutputTokens); - params.draftIds = bufferCast(*mDraftTokens); - params.targetIds = bufferCast(*mTargetTokens); - params.sequenceLengths = bufferCast(*mSequenceLengths); - params.acceptedLengths = bufferCast(*mAcceptedLengths); - params.finishedFinal - = reinterpret_cast(bufferCast(*mFinishedFinal)); - params.batchSlots = bufferCast(*mBatchSlots); - params.paths = bufferCast(*mPaths); - params.endIds = bufferCast(*mEndIds); - params.medusaLogits = reinterpret_cast(bufferCast(*mMedusaInputLogitsPtrs)); - params.logitsPtrs = reinterpret_cast(bufferCast(*mMedusaLogitsPtrs)); - params.curTokensPerStep = bufferCast(*mTokensPerStep); - params.targetTokensPerStep = bufferCast(*mTokensPerStep); - params.bestPathIds = bufferCast(*mBestPaths); - params.batchSize = mBatchSize; - params.maxBatchSize = mMaxBatchSize; - params.vocabSize = mVocabSize; - params.maxSeqLen = mMaxSeqLen; - params.maxDraftPathLen = mMaxNumHeads; - params.maxDecodingTokens = mMaxDraftSeqPerStep; - params.stream = mStream->get(); - - params.checkParams(); - - tksp::acceptDraftTokensByIdsWithPaths(params); - } - - void callTestedKernel() - { - switch (mAcceptMode) - { - case AcceptKernelMode::BY_IDS: callAcceptByIds(); break; - case AcceptKernelMode::BY_LOGITS: callAcceptByLogits(); break; - case AcceptKernelMode::BY_IDS_WITH_PATH: callAcceptByIdsWithPaths(); break; - default: TLLM_CHECK(false); // Should never be here - } - } - - void verifyAcceptByIdsResults(SizeType32 seed) - { - auto finishedFinalPtr - = reinterpret_cast(bufferCast(*mFinishedFinal)); - auto sequenceLengthsPtr = BufferRange(*mSequenceLengths); - auto finishedSumPtr = BufferRange(*mFinishedSum); - auto batchSlotsPtr = BufferRange(*mBatchSlots); - // Verify seqLen for accepted tokens - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const batchSlot = batchSlotsPtr[bi]; - EXPECT_EQ(mOutputLen[batchSlot], sequenceLengthsPtr[batchSlot]) << " bi " << bi << " seed " << seed; - EXPECT_EQ(mAcceptedFinished[batchSlot].isFinished(), finishedFinalPtr[batchSlot].isFinished()) - << " bi " << bi << " seed " << seed; - EXPECT_EQ(mAcceptedFinished[batchSlot].isSkipDecoding(), finishedFinalPtr[batchSlot].isSkipDecoding()) - << " bi " << bi << " seed " << seed; - EXPECT_EQ(static_cast(mAcceptedFinished[batchSlot].isFinished()), finishedSumPtr[batchSlot]); - } - } - - void verifyAcceptByLogitsResults(SizeType32 seed) - { - auto finishedStepsPtr - = reinterpret_cast(bufferCast(*mFinishedSteps)); - auto contextLengthsPtr = BufferRange(*mContextLengths); - auto outLogitsPtr = BufferRange(*mTargetLogits); - auto refLogitsPtr = BufferRange(*mRefTargetLogits); - auto numsDraftTokensPtr = BufferRange(*mNumsDraftTokens); - auto batchSlotsPtr = BufferRange(*mBatchSlots); - - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const batchSlot = batchSlotsPtr[bi]; - for (SizeType32 si = 0; si < numsDraftTokensPtr[batchSlot]; ++si) - { - auto const outFinishedState = finishedStepsPtr[si * mMaxBatchSize + batchSlot]; - auto const logitsOffset = batchSlot * mMaxDraftTokens * mVocabSize + si * mVocabSize; - if (si <= mAcceptedLen[batchSlot] - contextLengthsPtr[batchSlot]) - { - EXPECT_FALSE(outFinishedState.isSkipDecoding()) - << " bi: " << bi << " si: " << si << " seed: " << seed; - for (SizeType32 vi = 0; vi < mVocabSize; ++vi) - { - auto const outLogit = static_cast(outLogitsPtr[logitsOffset + vi]); - auto const refLogit = static_cast(refLogitsPtr[logitsOffset + vi]); - EXPECT_FALSE((refLogit > -10) ^ (outLogit > -10)) - << " bi: " << bi << " si: " << si << " vi: " << vi << " seed: " << seed; - if (refLogit > -10 && outLogit > -10) - { - if (!almostEqual(outLogit, refLogit, 1e-1, 1e-2)) - { - std::cout << refLogit << " " << outLogit << std::endl; - } - ASSERT_TRUE(almostEqual(outLogit, refLogit, 1e-1, 1e-2)) - << " bi: " << bi << " si: " << si << " vi: " << vi << " seed: " << seed; - } - } - } - else - { - EXPECT_TRUE(outFinishedState.isSkipDecoding()) - << " bi: " << bi << " si: " << si << " seed: " << seed; - } - } - } - } - - void verifyAcceptByIdsWithPathsResults(SizeType32 seed) - { - auto medusaLogitsPtrsPtr = BufferRange(*mMedusaLogitsPtrs); - auto batchSlotsPtr = BufferRange(*mBatchSlots); - auto draftContextLengths = BufferRange(*mSequenceLengths); - auto draftContextLengthsInit = BufferRange(*mSequenceLengthsCopy); - auto acceptedLengths = BufferRange(*mAcceptedLengths); - auto outputIdsPtr = BufferRange(*mOutputTokens); - auto bestPathIds = BufferRange(*mBestPaths); - auto finishedFinalPtr - = reinterpret_cast(bufferCast(*mFinishedFinal)); - - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const batchSlot = batchSlotsPtr[bi]; - auto const bestPathIdx = mAcceptedPathIdx[batchSlot]; - auto const lastTargetIdx = mLastTargetIdx[batchSlot]; - if (lastTargetIdx < 0) - { - continue; - } - - auto const acceptedLen = mAcceptedLen[batchSlot]; - auto acceptedTokens = mRefAcceptedTokens[batchSlot]; - - EXPECT_EQ(bestPathIds[batchSlot], bestPathIdx) << "bi: " << bi << " seed: " << seed; - - for (int32_t hi = 0; hi < mMaxNumHeads; ++hi) - { - auto refOffset - = tc::flat_index4(hi, batchSlot, lastTargetIdx, 0, mMaxBatchSize, mMaxDraftSeqPerStep, mVocabSize); - auto outOffset - = static_cast(medusaLogitsPtrsPtr[bi * mMaxNumHeads + hi] - static_cast(nullptr)); - EXPECT_EQ(outOffset, refOffset) << " bi: " << bi << " hi: " << hi << " seed: " << seed; - } - EXPECT_EQ(acceptedLengths[batchSlot], acceptedLen) << " bi: " << bi << " seed: " << seed; - EXPECT_EQ(draftContextLengths[batchSlot], draftContextLengthsInit[batchSlot] + acceptedLen) - << " bi: " << bi << " seed: " << seed << " out: " << draftContextLengths[batchSlot] - << " ref: " << draftContextLengthsInit[batchSlot] + acceptedLen; - - for (SizeType32 ti = 0; ti < acceptedLen; ++ti) - { - ASSERT_EQ(mRefAcceptedTokens[batchSlot].size(), acceptedLen) - << " bi: " << bi << " ti: " << ti << " seed: " << seed; - EXPECT_EQ(outputIdsPtr[batchSlot * mMaxSeqLen + draftContextLengthsInit[batchSlot] + ti], - mRefAcceptedTokens[batchSlot][ti]) - << " bi: " << bi << " ti: " << ti << " seed: " << seed; - } - EXPECT_EQ(finishedFinalPtr[batchSlot].isFinished(), mFinishedByIdsPaths[batchSlot]) - << " bi: " << bi << " seed: " << seed; - } - } - - void verifyResult(SizeType32 seed) - { - switch (mAcceptMode) - { - case AcceptKernelMode::BY_IDS: verifyAcceptByIdsResults(seed); break; - case AcceptKernelMode::BY_LOGITS: verifyAcceptByLogitsResults(seed); break; - case AcceptKernelMode::BY_IDS_WITH_PATH: verifyAcceptByIdsWithPathsResults(seed); break; - default: TLLM_CHECK(false); // Should never be here - } - } - - void runTest(DecodingKernelTestParam const& params) - { - mAcceptMode = params.mAcceptMode; - - mBatchSize = params.mBatchSize; - mMaxBatchSize = params.mMaxBatchSize; - mBeamWidth = params.mBeamWidth; - mVocabSize = params.mVocabSize; - mMaxDraftTokens = params.mMaxDraftTokens; - mMaxSeqLen = params.mMaxSeqLen; - - mMaxNumHeads = params.mMaxNumHeads; - if (mMaxNumHeads > 1 && mAcceptMode != AcceptKernelMode::BY_IDS_WITH_PATH) - { - GTEST_SKIP() << "MaxNumHeads > 1 is only supported for AcceptKernelMode::BY_IDS_WITH_PATH"; - } - - mMaxDraftSeqPerStep = params.mMaxDraftSeqPerStep; - if (mMaxDraftSeqPerStep > 1 && mAcceptMode != AcceptKernelMode::BY_IDS_WITH_PATH) - { - GTEST_SKIP() << "MaxDraftSeqPerStep > 1 is only supported for AcceptKernelMode::BY_IDS_WITH_PATH"; - } - - mMaxTotalDraftTokens = mMaxDraftSeqPerStep * mMaxDraftTokens; - mPadId = mVocabSize - 1; - - mMaxDraftSeqlen = mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH ? mMaxDraftTokens - 1 : mMaxDraftTokens; - mMaxTargetSeqlen = mAcceptMode == AcceptKernelMode::BY_IDS_WITH_PATH ? mMaxDraftTokens : mMaxSeqLen; - - createBuffers(); - - for (SizeType32 seed = 0; seed < mSeeds; ++seed) - { - // if (seed != 145) - // { - // continue; - // } - TLLM_LOG_DEBUG("Seed %d", seed); - - initData(seed); - - mStream->synchronize(); - - callTestedKernel(); - - mStream->synchronize(); - - verifyResult(seed); - } - } - -protected: - std::shared_ptr mBufferManager; - std::shared_ptr mStream; - - TensorPtr mDraftTokens; - TensorPtr mTargetTokens; - TensorPtr mOutputTokens; - - TensorPtr mDraftLogits; - TensorPtr mTargetLogits; - TensorPtr mTargetLogitsPtrs; - TensorPtr mRefTargetLogits; - - TensorPtr mDraftProbs; - TensorPtr mTargetProbs; - - TensorPtr mNumsDraftTokens; - TensorPtr mSequenceLengths; - TensorPtr mSequenceLengthsCopy; - TensorPtr mAcceptedLengths; - TensorPtr mContextLengths; - TensorPtr mFinishedSteps; - TensorPtr mFinishedFinal; - TensorPtr mFinishedSum; - TensorPtr mBatchSlots; - - TensorPtr mPaths; - TensorPtr mEndIds; - TensorPtr mMedusaLogitsPtrs; - TensorPtr mMedusaInputLogitsPtrs; - TensorPtr mTokensPerStep; - TensorPtr mBestPaths; - - TensorPtr mCurandStates; - - std::vector mAcceptedLen; - std::vector mOutputLen; - std::vector mAcceptedFinished; - std::vector mAcceptedPathIdx; - std::vector mLastTargetIdx; - std::vector> mRefAcceptedTokens; - std::vector mFinishedByIdsPaths; - - SizeType32 mBatchSize; - SizeType32 mMaxBatchSize; - SizeType32 mBeamWidth; - SizeType32 mMaxSeqLen; - SizeType32 mVocabSize; - SizeType32 mMaxDraftTokens; - SizeType32 mMaxTotalDraftTokens; - SizeType32 mMaxDraftSeqlen; - SizeType32 mMaxTargetSeqlen; - SizeType32 mMaxNumHeads; - SizeType32 mMaxDraftSeqPerStep; - AcceptKernelMode mAcceptMode; - SizeType32 mPadId; - static constexpr SizeType32 mSeeds = 64; -}; - -template class DecodingKernelsTest; -template class DecodingKernelsTest; - -typedef testing::Types FloatAndHalfTypes; - -TYPED_TEST_SUITE(DecodingKernelsTest, FloatAndHalfTypes); - -TYPED_TEST(DecodingKernelsTest, DISABLED_acceptDraftTokensByIdsKernelSmall) -{ - this->runTest(DecodingKernelTestParam() - .setBatchSize(1) - .setMaxSeqLen(16) - .setVocabSize(32) - .setMaxDraftTokens(8) - .setMaxDraftSeqPerStep(1) - .setAcceptMode(AcceptKernelMode::BY_IDS)); -} - -TYPED_TEST(DecodingKernelsTest, DISABLED_acceptDraftTokensByIdsKernelLarge) -{ - this->runTest(DecodingKernelTestParam() - .setBatchSize(128) - .setMaxSeqLen(128) - .setVocabSize(52000) - .setMaxDraftTokens(8) - .setMaxDraftSeqPerStep(1) - .setAcceptMode(AcceptKernelMode::BY_IDS)); -} - -TYPED_TEST(DecodingKernelsTest, DISABLED_acceptDraftTokensByLogitsKernelSmall) -{ - this->runTest(DecodingKernelTestParam() - .setBatchSize(1) - .setMaxSeqLen(16) - .setVocabSize(32) - .setMaxDraftTokens(8) - .setMaxDraftSeqPerStep(1) - .setAcceptMode(AcceptKernelMode::BY_LOGITS)); -} - -TYPED_TEST(DecodingKernelsTest, DISABLED_acceptDraftTokensByLogitsKernelLarge) -{ - this->runTest(DecodingKernelTestParam() - .setBatchSize(64) - .setMaxSeqLen(64) - .setVocabSize(4000) - .setMaxDraftTokens(8) - .setMaxDraftSeqPerStep(1) - .setAcceptMode(AcceptKernelMode::BY_LOGITS)); -} - -// FIXME: test is incorrect and too complicated. -TYPED_TEST(DecodingKernelsTest, DISABLED_acceptDraftTokensByIdsWithPathsKernelSmall) -{ - this->runTest(DecodingKernelTestParam() - .setBatchSize(1) - .setMaxSeqLen(128) - .setVocabSize(32) - .setMaxDraftTokens(3) - .setMaxDraftSeqPerStep(4) - .setMaxNumHeads(2) - .setAcceptMode(AcceptKernelMode::BY_IDS_WITH_PATH)); -} - -// FIXME: test is incorrect and too complicated. -TYPED_TEST(DecodingKernelsTest, DISABLED_acceptDraftTokensByIdsWithPathsKernelLarge) -{ - this->runTest(DecodingKernelTestParam() - .setBatchSize(128) - .setMaxSeqLen(1024) - .setVocabSize(4000) - .setMaxDraftTokens(8) - .setMaxDraftSeqPerStep(64) - .setMaxNumHeads(7) - .setAcceptMode(AcceptKernelMode::BY_IDS_WITH_PATH)); -} -} // end of namespace diff --git a/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp b/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp deleted file mode 100644 index 8ce24f813e65..000000000000 --- a/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp +++ /dev/null @@ -1,518 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include - -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -#include "tensorrt_llm/common/tllmDataType.h" - -#include -#include -#include - -namespace -{ - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::common; - -namespace tk = tensorrt_llm::kernels; -namespace trk = tensorrt_llm::runtime::kernels; -namespace tksd = tensorrt_llm::kernels::speculative_decoding; - -class SamplingParams -{ -public: - SamplingParams() {} - - inline void setNumCtxRequests(SizeType32 numCtxRequests) - { - mNumCtxRequests = numCtxRequests; - } - - inline void setNumGenRequests(SizeType32 numGenRequests) - { - mNumGenRequests = numGenRequests; - } - - inline void setMaxPathLen(SizeType32 maxPathLen) - { - mMaxPathLen = maxPathLen; - } - - [[nodiscard]] inline SizeType32 getNumCtxRequests() const - { - return mNumCtxRequests; - } - - [[nodiscard]] inline SizeType32 getNumGenRequests() const - { - return mNumGenRequests; - } - - [[nodiscard]] inline SizeType32 getBatchSize() const - { - return getNumCtxRequests() + getNumGenRequests(); - } - - [[nodiscard]] inline SizeType32 getVocabSize() const - { - return mVocabSize; - } - - [[nodiscard]] inline SizeType32 getMaxBatchSize() const - { - return 2 * getBatchSize(); - } - - [[nodiscard]] inline SizeType32 getMaxPathLen() const - { - return mMaxPathLen; - } - - [[nodiscard]] inline SizeType32 getMaxDecodingTokens() const - { - return mMaxDecodingTokens; - } - - [[nodiscard]] inline SizeType32 getMaxDecodingDraftTokens() const - { - return getMaxDecodingTokens() - 1; - } - - [[nodiscard]] inline SizeType32 getMaxSeqLen() const - { - return getMaxDecodingTokens() * 2; - } - -private: - SizeType32 mNumCtxRequests{6}; - SizeType32 mNumGenRequests{6}; - SizeType32 mMaxPathLen{4}; - SizeType32 mMaxDecodingTokens{32}; - SizeType32 mVocabSize{256}; -}; - -class EaglePackDataTest : public ::testing::Test -{ -public: - using BufferPtr = IBuffer::SharedPtr; - using TensorPtr = ITensor::SharedPtr; - - void SetUp() override - { - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); - } - - void allocateBuffers() - { - // inputs - mBatchSlots = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mInputTemperatures = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); - - mInputRandomDataSample = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); - - mInputRandomDataValidation = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kFLOAT); - - mInputNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mInputNextDraftPaths - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize(), - mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mInputSpecDecodingGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mInputSpecDecodingPositionOffsets = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - - auto const numPackedMasks - = static_cast(tensorrt_llm::common::divUp(mSamplingParams.getMaxDecodingTokens(), 32)); - mInputSpecDecodingPackedMasks = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), numPackedMasks}), - tensorrt_llm::DataType::kINT32); - - // outputs - mOutputTemperatures = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kFLOAT); - - mOutputRandomDataSample = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kFLOAT); - - mOutputRandomDataValidation = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kFLOAT); - - mOutputNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mOutputNextDraftLens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mOutputNextDraftPaths - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize(), - mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mOutputSpecDecodingGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mOutputSpecDecodingPositionOffsets = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - - mOutputSpecDecodingPackedMasks = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), numPackedMasks}), - tensorrt_llm::DataType::kINT32); - - // workspace - mMaxGenerationLength = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - mCumSumGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize() + 1}), tensorrt_llm::DataType::kINT32); - - mScanReduceTempStorageBytes = tksd::invokeScanReduceGenerationLengths( - mSamplingParams.getBatchSize(), nullptr, nullptr, 0, nullptr, nullptr, mStream->get()); - mScanReduceTempStorage = mBufferManager->gpu(mScanReduceTempStorageBytes); - } - - void initBuffers() - { - trk::invokeFill(*mOutputTemperatures, float{0}, *mStream); - trk::invokeFill(*mOutputRandomDataSample, float{0}, *mStream); - trk::invokeFill(*mOutputRandomDataValidation, float{0}, *mStream); - trk::invokeFill(*mOutputNextDraftTokens, TokenIdType{-1}, *mStream); - trk::invokeFill(*mOutputNextDraftLens, SizeType32{0}, *mStream); - trk::invokeFill(*mOutputNextDraftPaths, SizeType32{0}, *mStream); - trk::invokeFill(*mOutputSpecDecodingGenerationLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mOutputSpecDecodingPositionOffsets, SizeType32{0}, *mStream); - trk::invokeFill(*mOutputSpecDecodingPackedMasks, SizeType32{0}, *mStream); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - batchSlotsPtr[bi] = 2 * bi; - } - - std::mt19937 gen(42); - std::uniform_real_distribution distr(0.0, 1.0); - std::uniform_int_distribution intDistr(0, 1000); - std::uniform_int_distribution lenDistr(0, mSamplingParams.getMaxDecodingTokens() - 1); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - bufferCast(*mInputTemperatures)[batchSlotsPtr[bi]] = distr(gen); - bufferCast(*mInputRandomDataSample)[batchSlotsPtr[bi]] = distr(gen); - } - - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDecodingDraftTokens(); ++ti) - { - bufferCast(*mInputNextDraftTokens)[flat_index2( - batchSlotsPtr[bi], ti, mSamplingParams.getMaxDecodingDraftTokens())] - = intDistr(gen); - } - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDecodingTokens(); ++ti) - { - bufferCast( - *mInputRandomDataValidation)[batchSlotsPtr[bi] * mSamplingParams.getMaxDecodingTokens() + ti] - = distr(gen); - for (SizeType32 pi = 0; pi < mSamplingParams.getMaxPathLen(); ++pi) - { - bufferCast(*mInputNextDraftPaths)[flat_index3(batchSlotsPtr[bi], ti, pi, - mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen())] - = intDistr(gen); - } - auto const numPackedMasks - = static_cast(tensorrt_llm::common::divUp(mSamplingParams.getMaxDecodingTokens(), 32)); - for (SizeType32 mi = 0; mi < numPackedMasks; ++mi) - { - bufferCast(*mInputSpecDecodingPackedMasks)[flat_index3( - batchSlotsPtr[bi], ti, mi, mSamplingParams.getMaxDecodingTokens(), numPackedMasks)] - = intDistr(gen); - } - bufferCast(*mInputSpecDecodingPositionOffsets)[flat_index2( - batchSlotsPtr[bi], ti, mSamplingParams.getMaxDecodingTokens())] - = intDistr(gen); - } - bufferCast(*mInputSpecDecodingGenerationLengths)[batchSlotsPtr[bi]] = lenDistr(gen) + 1; - } - } - - void callPackData() - { - tksd::PackEagleParams params; - params.batchSize = mSamplingParams.getBatchSize(); - params.maxNumPaths = mSamplingParams.getMaxDecodingTokens(); - params.maxDecodingTokens = mSamplingParams.getMaxDecodingTokens(); - params.maxPathLength = mSamplingParams.getMaxPathLen(); - params.numContextRequests = mSamplingParams.getNumCtxRequests(); - params.numGenerationRequests = mSamplingParams.getNumGenRequests(); - - params.batchSlots = bufferCast(*mBatchSlots); - - // Outputs from decoder -- inputs to the packing kernel - params.inputTemperatures = bufferCast(*mInputTemperatures); - params.inputRandomDataSample = bufferCast(*mInputRandomDataSample); - params.inputRandomDataValidation = bufferCast(*mInputRandomDataValidation); - - params.inputNextDraftTokens = bufferCast(*mInputNextDraftTokens); - params.inputNextDraftPaths = bufferCast(*mInputNextDraftPaths); - - params.inputSpecDecodingGenerationLengths = bufferCast(*mInputSpecDecodingGenerationLengths); - params.inputSpecDecodingPositionOffsets = bufferCast(*mInputSpecDecodingPositionOffsets); - params.inputSpecDecodingPackedMasks = bufferCast(*mInputSpecDecodingPackedMasks); - - // Outputs of the packing kernel -- inputs to the engine - params.outputTemperatures = bufferCast(*mOutputTemperatures); - params.outputRandomDataSample = bufferCast(*mOutputRandomDataSample); - params.outputRandomDataValidation = bufferCast(*mOutputRandomDataValidation); - - params.outputNextDraftTokens = bufferCast(*mOutputNextDraftTokens); - params.outputNextDraftLens = bufferCast(*mOutputNextDraftLens); - params.outputNextDraftPaths = bufferCast(*mOutputNextDraftPaths); - - params.outputSpecDecodingGenerationLengths = bufferCast(*mOutputSpecDecodingGenerationLengths); - params.outputSpecDecodingPositionOffsets = bufferCast(*mOutputSpecDecodingPositionOffsets); - params.outputSpecDecodingPackedMasks = bufferCast(*mOutputSpecDecodingPackedMasks); - - params.maxGenerationLength = bufferCast(*mMaxGenerationLength); - params.cumSumGenerationLengths = bufferCast(*mCumSumGenerationLengths); - - params.checkParams(); - - if (mSamplingParams.getNumGenRequests()) - { - // Pack tensors from batch slot position to continuous array - tksd::invokePackEagleGenerationLengths(params, mStream->get()); - - sync_check_cuda_error(mStream->get()); - - // Compute inclusive sum and max - tksd::invokeScanReduceGenerationLengths(mSamplingParams.getNumGenRequests(), - bufferCast(*mOutputSpecDecodingGenerationLengths), - bufferCast(*mScanReduceTempStorage), mScanReduceTempStorageBytes, - bufferCast(*mCumSumGenerationLengths), bufferCast(*mMaxGenerationLength), - mStream->get()); - - sync_check_cuda_error(mStream->get()); - } - - mStream->synchronize(); - - // Pack tensors from batch slot position to continuous array - tksd::invokePackEagle(params, mStream->get()); - - sync_check_cuda_error(mStream->get()); - } - - void verifyResults() - { - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - EXPECT_EQ(BufferRange(*mInputTemperatures)[batchSlotsPtr[bi]], - BufferRange(*mOutputTemperatures)[bi]); - EXPECT_EQ(BufferRange(*mInputRandomDataSample)[batchSlotsPtr[bi]], - BufferRange(*mOutputRandomDataSample)[bi]); - } - - auto const numCtxRequests = mSamplingParams.getNumCtxRequests(); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDecodingTokens(); ++ti) - { - EXPECT_EQ( - BufferRange( - *mInputRandomDataValidation)[batchSlotsPtr[bi] * mSamplingParams.getMaxDecodingTokens() + ti], - BufferRange(*mOutputRandomDataValidation)[bi * mSamplingParams.getMaxDecodingTokens() + ti]); - for (SizeType32 pi = 0; pi < mSamplingParams.getMaxPathLen(); ++pi) - { - EXPECT_EQ(BufferRange(*mInputNextDraftPaths)[flat_index3(batchSlotsPtr[bi], ti, pi, - mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen())], - BufferRange(*mOutputNextDraftPaths)[flat_index3( - bi, ti, pi, mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen())]); - } - } - EXPECT_EQ(BufferRange(*mOutputNextDraftLens)[bi], - bi < numCtxRequests - ? 0 - : BufferRange(*mInputSpecDecodingGenerationLengths)[batchSlotsPtr[bi]] - 1); - } - - auto const maxGenerationLength = bufferCast(*mMaxGenerationLength)[0]; - for (SizeType32 bi = 0; bi < mSamplingParams.getNumGenRequests(); ++bi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDecodingDraftTokens(); ++ti) - { - EXPECT_EQ(BufferRange(*mInputNextDraftTokens)[flat_index2( - batchSlotsPtr[numCtxRequests + bi], ti, mSamplingParams.getMaxDecodingDraftTokens())], - BufferRange(*mOutputNextDraftTokens)[flat_index2( - numCtxRequests + bi, ti, mSamplingParams.getMaxDecodingDraftTokens())]); - } - EXPECT_EQ(BufferRange(*mInputSpecDecodingGenerationLengths)[batchSlotsPtr[numCtxRequests + bi]], - BufferRange(*mOutputSpecDecodingGenerationLengths)[bi]); - for (SizeType32 ti = 0; ti < maxGenerationLength; ++ti) - { - EXPECT_EQ(BufferRange(*mInputSpecDecodingPositionOffsets)[flat_index2( - batchSlotsPtr[numCtxRequests + bi], ti, mSamplingParams.getMaxDecodingTokens())], - BufferRange( - *mOutputSpecDecodingPositionOffsets)[flat_index2(bi, ti, maxGenerationLength)]) - << "bi: " << bi << " ti: " << ti; - } - auto const numTokens = (bi == 0) ? bufferCast(*mCumSumGenerationLengths)[0] - : bufferCast(*mCumSumGenerationLengths)[bi] - - bufferCast(*mCumSumGenerationLengths)[bi - 1]; - auto const outputStartId = (bi == 0) ? 0 : bufferCast(*mCumSumGenerationLengths)[bi - 1]; - auto const numPackedMasks - = static_cast(tensorrt_llm::common::divUp(mSamplingParams.getMaxDecodingTokens(), 32)); - for (SizeType32 ti = 0; ti < numTokens * numPackedMasks; ++ti) - { - EXPECT_EQ(BufferRange( - *mInputSpecDecodingPackedMasks)[flat_index2(batchSlotsPtr[numCtxRequests + bi], ti, - mSamplingParams.getMaxDecodingTokens() * numPackedMasks)], - BufferRange( - *mOutputSpecDecodingPackedMasks)[flat_index2(outputStartId, ti, numPackedMasks)]) - << "bi: " << bi << " ti: " << ti; - } - } - } - - void run(SamplingParams samplingParams) - { - mSamplingParams = samplingParams; - - allocateBuffers(); - - initBuffers(); - - callPackData(); - - mStream->synchronize(); - - verifyResults(); - } - -private: - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - - // input - TensorPtr mBatchSlots; - TensorPtr mInputTemperatures; - TensorPtr mInputRandomDataSample; - TensorPtr mInputRandomDataValidation; - TensorPtr mInputNextDraftTokens; - TensorPtr mInputNextDraftPaths; - TensorPtr mInputSpecDecodingGenerationLengths; - TensorPtr mInputSpecDecodingPositionOffsets; - TensorPtr mInputSpecDecodingPackedMasks; - - // output - TensorPtr mOutputTemperatures; - TensorPtr mOutputRandomDataSample; - TensorPtr mOutputRandomDataValidation; - TensorPtr mOutputNextDraftTokens; - TensorPtr mOutputNextDraftLens; - TensorPtr mOutputNextDraftPaths; - TensorPtr mOutputSpecDecodingGenerationLengths; - TensorPtr mOutputSpecDecodingPositionOffsets; - TensorPtr mOutputSpecDecodingPackedMasks; - - // workspace - TensorPtr mMaxGenerationLength; - TensorPtr mCumSumGenerationLengths; - - BufferPtr mScanReduceTempStorage; - - SizeType32 mScanReduceTempStorageBytes; - - SamplingParams mSamplingParams; -}; - -TEST_F(EaglePackDataTest, Ctx1Gen0) -{ - SamplingParams params; - - params.setNumCtxRequests(1); - params.setNumGenRequests(0); - - this->run(params); -} - -TEST_F(EaglePackDataTest, Ctx0Gen1) -{ - SamplingParams params; - - params.setNumCtxRequests(0); - params.setNumGenRequests(1); - - this->run(params); -} - -TEST_F(EaglePackDataTest, Ctx100Gen0) -{ - SamplingParams params; - - params.setNumCtxRequests(100); - params.setNumGenRequests(0); - - this->run(params); -} - -TEST_F(EaglePackDataTest, Ctx0Gen100) -{ - SamplingParams params; - - params.setNumCtxRequests(0); - params.setNumGenRequests(100); - - this->run(params); -} - -TEST_F(EaglePackDataTest, Ctx100Gen100) -{ - SamplingParams params; - - params.setNumCtxRequests(100); - params.setNumGenRequests(100); - - this->run(params); -} -} // namespace diff --git a/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu b/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu index f0f7690257e7..e553abae39ed 100644 --- a/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu +++ b/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu @@ -15,8 +15,8 @@ */ #include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/runtime/gptDecoder.h" #include "tests/unit_tests/kernels/sampling/samplingTest.h" +#include #include using namespace tensorrt_llm::tests::kernels::sampling; @@ -28,6 +28,15 @@ namespace tk = tensorrt_llm::kernels; namespace { +/// @brief Helper function to produce batch slots [0, 1, ..., batchSize - 1]. +ITensor::SharedConstPtr getDefaultBatchSlots(SizeType32 batchSize) +{ + auto defaultBatchSlots = BufferManager::pinnedPool(ITensor::makeShape({batchSize}), TRTDataType::value); + auto range = BufferRange(*defaultBatchSlots); + std::iota(range.begin(), range.end(), 0); + return defaultBatchSlots; +} + __global__ void generateRandomNumber( SizeType32* vals, SizeType32 const* batchSlots, curandState_t* states, SizeType32 batchSize) { diff --git a/cpp/tests/unit_tests/layers/CMakeLists.txt b/cpp/tests/unit_tests/layers/CMakeLists.txt deleted file mode 100644 index 3f194d3a77c9..000000000000 --- a/cpp/tests/unit_tests/layers/CMakeLists.txt +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & -# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); you may not -# use this file except in compliance with the License. You may obtain a copy of -# the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT -# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the -# License for the specific language governing permissions and limitations under -# the License. - -set(SAMPLING_LAYER_TEST_SRC - baseSamplingLayerTest.cpp samplingLayerTest.cpp topKSamplingLayerTest.cpp - topPSamplingLayerTest.cpp externalDraftTokensLayerTest.cpp) -add_gtest(samplingLayerTest "${SAMPLING_LAYER_TEST_SRC}") - -set(BEAM_SEARCH_LAYER_TEST_SRC baseSamplingLayerTest.cpp - beamSearchLayerTest.cpp) -add_gtest(beamSearchLayerTest "${BEAM_SEARCH_LAYER_TEST_SRC}") - -set(LOOKAHEAD_POOLMANAGER_TEST_SRC randomLlm.cpp lookaheadPoolManagerTest.cpp) -add_gtest(lookaheadPoolManagerTest "${LOOKAHEAD_POOLMANAGER_TEST_SRC}") - -set(LOOKAHEAD_ALGORITHM_TEST_SRC randomLlm.cpp lookaheadAlgorithmTest.cpp) -add_gtest(lookaheadAlgorithmTest "${LOOKAHEAD_ALGORITHM_TEST_SRC}") - -set(LOOKAHEAD_RANDOMLLM_TEST_SRC randomLlm.cpp lookaheadRandomLlmTest.cpp) -add_gtest(lookaheadRandomLlmTest "${LOOKAHEAD_RANDOMLLM_TEST_SRC}") - -set(LOOKAHEAD_DECODING_TEST_SRC randomLlm.cpp lookaheadDecodingLayerTest.cpp) -add_gtest(lookaheadDecodingLayerTest "${LOOKAHEAD_DECODING_TEST_SRC}") - -add_gtest(dynamicDecodeLayerTest dynamicDecodeLayerTest.cpp) -add_gtest(eagleLayerTest eagleLayerTest.cpp) -add_gtest(explicitDraftTokensLayerTest explicitDraftTokensLayerTest.cpp) -add_gtest(layerUtilsTest layerUtilsTest.cpp) -add_gtest(medusaDecodeLayerTest medusaDecodeLayerTest.cpp) diff --git a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp b/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp deleted file mode 100644 index 7886a6b54e2b..000000000000 --- a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp +++ /dev/null @@ -1,375 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tests/unit_tests/layers/baseSamplingLayerTest.h" -#include "tensorrt_llm/common/tllmDataType.h" - -namespace tensorrt_llm::tests::layers::sampling -{ - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using namespace tensorrt_llm::common; - -namespace tk = tensorrt_llm::kernels; -namespace trk = tensorrt_llm::runtime::kernels; - -template -void BaseSamplingLayerTest::setup(uint64_t seed, TestSamplingParams const& params) -{ - auto const dataType = TRTDataType::value; - auto const ptrType = TRTDataType::value; - - // clang-format off - - // logits = (-0.9163, -1.2040, -1.6094, -2.3026) -> prob = (0.4, 0.3, 0.2, 0.1) - std::vector testLogits = { - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, // step 0 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // step 1 - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // step 2 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX // step 3 - }; - - // clang-format on - - if (params.beamWidth == 1) - { - mTestLogitsInit = testLogits; - } - else - { - for (int step = 0; step < mMaxSeqLen; ++step) - { - auto const& logitsBegin = testLogits.begin() + mVocabSize * step; - auto const& logitsEnd = testLogits.begin() + mVocabSize * (step + 1); - for (int bm = 0; bm < params.beamWidth; ++bm) - { - mTestLogitsInit.insert(mTestLogitsInit.end(), logitsBegin, logitsEnd); - } - } - } - - if (mComputeProbs) - { - computeProb(mTestLogitsInit.data(), mTestLogitsInit.data(), - BaseSamplingLayerTest::mMaxOutputLen * params.beamWidth, mVocabSize); - } - - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kINT32); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kINT32); - mFinishedDevice = params.isExternalDraftTokensLayerTest - ? mBufferManager->gpu(ITensor::makeShape({mMaxTokensPerEngineStep, maxBatchSize()}), - TRTDataType::value) - : mBufferManager->gpu( - ITensor::makeShape({maxBatchSize()}), TRTDataType::value); - mOutputIdsDevice = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kINT32); - mIdsPtrHost = mBufferManager->pinned(ITensor::makeShape({maxBatchSize()}), ptrType); - - mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kFLOAT); - mOutputLogProbsDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), mMaxSeqLen}), tensorrt_llm::DataType::kFLOAT); - - mBatchSlots - = mBufferManager->pinned(ITensor::makeShape({mBatchSize + mBatchSizeBadPad}), tensorrt_llm::DataType::kINT32); - mCurandStatesDevice = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); - - auto const workspaceSize = mSamplingLayer->getWorkspaceSize(); - - trk::invokeFill(*mSeqLengthsDevice, int32_t{0}, *mStream); - trk::invokeFill(*mContextLengthDevice, int32_t{0}, *mStream); - trk::invokeFill(*mFinishedDevice, uint8_t{0}, *mStream); - trk::invokeFill(*mOutputIdsDevice, int32_t{0}, *mStream); - trk::invokeFill(*mCumLogProbsDevice, float{0.0f}, *mStream); - trk::invokeFill(*mOutputLogProbsDevice, float{0.0f}, *mStream); - trk::invokeFill(*mEndIdsDevice, int32_t{mEndId}, *mStream); - tk::invokeCurandInitialize(reinterpret_cast(bufferCast(*mCurandStatesDevice)), nullptr, - maxBatchSize(), seed, mStream->get()); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - batchSlotsPtr[bi] = kDoubleBatchIdx * bi; - } - for (SizeType32 bi = 0; bi < mBatchSizeBadPad; ++bi) - { - batchSlotsPtr[mBatchSize + bi] = 0xbaadf00d; - } - - auto idsPtrHostPtr = BufferRange(*mIdsPtrHost); - auto outputIdsDevicePtr = bufferCast(*mOutputIdsDevice); - for (SizeType32 bi = 0; bi < maxBatchSize(); bi++) - { - idsPtrHostPtr[bi] = outputIdsDevicePtr + bi * mMaxSeqLen; - } - - std::shared_ptr setupParams; - if (params.isExternalDraftTokensLayerTest) - { - auto externalDraftTokensSetupParams = std::make_shared(); - externalDraftTokensSetupParams->randomSeed = std::make_optional>({seed}); - externalDraftTokensSetupParams->runtimeTopK - = params.topKs.size() ? std::make_optional>(params.topKs) : std::nullopt; - externalDraftTokensSetupParams->runtimeTopP - = params.topPs.size() ? std::make_optional>(params.topPs) : std::nullopt; - - setupParams = externalDraftTokensSetupParams; - } - else if (mBeamWidth == 1) - { - auto samplingSetupParams = std::make_shared(); - samplingSetupParams->randomSeed = std::make_optional>({seed}); - samplingSetupParams->runtimeTopK - = params.topKs.size() ? std::make_optional>(params.topKs) : std::nullopt; - samplingSetupParams->runtimeTopP - = params.topPs.size() ? std::make_optional>(params.topPs) : std::nullopt; - samplingSetupParams->topPDecay - = params.decay.size() ? std::make_optional>(params.decay) : std::nullopt; - samplingSetupParams->topPMin - = params.minTopP.size() ? std::make_optional>(params.minTopP) : std::nullopt; - samplingSetupParams->topPResetIds - = params.topPResetIds.size() ? std::make_optional>(params.topPResetIds) : std::nullopt; - - setupParams = samplingSetupParams; - } - else // Beam Search - { - auto samplingSetupParams = std::make_shared(); - setupParams = samplingSetupParams; - - mSrcCacheIndirection = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mTgtCacheIndirection = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mParentIds = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - - auto constexpr nvTokenIdType = TRTDataType::value; - auto constexpr nvSizeType = TRTDataType::value; - auto constexpr nvFloatType = TRTDataType::value; - auto constexpr nvBoolType = TRTDataType::value; - mOutputIdsCBA - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), 2 * mBeamWidth, mMaxSeqLen}), nvTokenIdType); - mLogProbsCBA - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), 2 * mBeamWidth, mMaxSeqLen}), nvFloatType); - mSequenceLengthsCBA = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), 2 * mBeamWidth}), nvSizeType); - mCumLogProbsCBA = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), 2 * mBeamWidth}), nvFloatType); - mNormedScoresCBA = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), 2 * mBeamWidth}), nvFloatType); - mNumBeamsCBA = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvSizeType); - mMinNormedScoresCBA = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvFloatType); - mBatchDones = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvBoolType); - mOutputIdsPtr = mBufferManager->pinned(ITensor::makeShape({maxBatchSize()}), ptrType); - mParentIdsPtr = mBufferManager->pinned(ITensor::makeShape({maxBatchSize()}), ptrType); - - trk::invokeFill(*mSrcCacheIndirection, int32_t{0}, *mStream); - trk::invokeFill(*mTgtCacheIndirection, int32_t{0}, *mStream); - trk::invokeFill(*mParentIds, int32_t{0}, *mStream); - trk::invokeFill(*mOutputIdsCBA, int32_t{0}, *mStream); - trk::invokeFill(*mLogProbsCBA, float{0}, *mStream); - trk::invokeFill(*mSequenceLengthsCBA, int32_t{0}, *mStream); - trk::invokeFill(*mCumLogProbsCBA, float{0}, *mStream); - trk::invokeFill(*mNormedScoresCBA, float{0}, *mStream); - trk::invokeFill(*mNumBeamsCBA, int32_t{0}, *mStream); - trk::invokeFill(*mMinNormedScoresCBA, float{0}, *mStream); - trk::invokeFill(*mBatchDones, bool{0}, *mStream); - - auto outputIdsPtr = bufferCast(*mOutputIdsPtr); - auto parentIdsPtr = bufferCast(*mParentIdsPtr); - for (SizeType32 bi = 0; bi < maxBatchSize(); bi++) - { - outputIdsPtr[bi] = outputIdsDevicePtr + bi * mMaxSeqLen; - parentIdsPtr[bi] = outputIdsDevicePtr + bi * mMaxSeqLen; - } - } - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mDecodingWorkspace->getDeviceRuntimeLogits()->reshape(ITensor::makeShape({mBatchSize, mBeamWidth, mVocabSize})); - mSamplingLayer->setup(mBatchSize, mBeamWidth, mBatchSlots, setupParams, mDecodingWorkspace); - mStream->synchronize(); -} - -template -std::shared_ptr BaseSamplingLayerTest::createInputTensors(int32_t step) -{ - constexpr int32_t ite = 0; - - auto decodeInputTensors = (mBeamWidth > 1) - ? std::make_shared(mEndIdsDevice, mBatchSlots, step, ite, mBatchSize) - : std::make_shared(mEndIdsDevice, mBatchSlots, step, ite, mBatchSize); - decodeInputTensors->logits = mDecodingWorkspace->getDeviceRuntimeLogits(); - decodeInputTensors->inputLengths = mContextLengthDevice; - decodeInputTensors->finished = mFinishedDevice; - - if (mBeamWidth > 1) - { - decodeInputTensors->srcCacheIndirection = mSrcCacheIndirection; - } - else - { - auto samplingInputTensors = std::dynamic_pointer_cast(decodeInputTensors); - samplingInputTensors->probsComputed = mComputeProbs; - samplingInputTensors->curandStates = reinterpret_cast(bufferCast(*mCurandStatesDevice)); - } - - return decodeInputTensors; -} - -template -std::shared_ptr BaseSamplingLayerTest::createOutputTensors() -{ - // TODO: check log probs and cum_log_probs - - auto decodeOutputs = (mBeamWidth > 1) ? std::make_shared(mOutputIdsDevice) - : std::make_shared(mOutputIdsDevice); - decodeOutputs->outputIdsPtr = mIdsPtrHost; - decodeOutputs->outputIdsPtrHost = mIdsPtrHost; - decodeOutputs->sequenceLength = mSeqLengthsDevice; - decodeOutputs->finished = mFinishedDevice; - decodeOutputs->outputLogProbs = mOutputLogProbsDevice; - decodeOutputs->cumLogProbs = mCumLogProbsDevice; - - if (mBeamWidth > 1) - { - auto beamSearchOutputs = std::dynamic_pointer_cast(decodeOutputs); - beamSearchOutputs->tgtCacheIndirection = mTgtCacheIndirection; - beamSearchOutputs->parentIds = mParentIds; - beamSearchOutputs->parentIdsPtr = mParentIdsPtr; - beamSearchOutputs->beamHypotheses = std::make_unique(); - beamSearchOutputs->beamHypotheses->outputIdsCBA = bufferCast(*mOutputIdsCBA); - beamSearchOutputs->beamHypotheses->logProbsCBA = bufferCast(*mLogProbsCBA); - beamSearchOutputs->beamHypotheses->sequenceLengthsCBA = bufferCast(*mSequenceLengthsCBA); - beamSearchOutputs->beamHypotheses->cumLogProbsCBA = bufferCast(*mCumLogProbsCBA); - beamSearchOutputs->beamHypotheses->normedScoresCBA = bufferCast(*mNormedScoresCBA); - beamSearchOutputs->beamHypotheses->numBeamsCBA = bufferCast(*mNumBeamsCBA); - beamSearchOutputs->beamHypotheses->minNormedScoresCBA = bufferCast(*mMinNormedScoresCBA); - beamSearchOutputs->beamHypotheses->batchDones = bufferCast(*mBatchDones); - } - - return decodeOutputs; -} - -template -void BaseSamplingLayerTest::batchCopy(int32_t step) -{ - auto const logitsHost = ITensor::wrap(mTestLogitsInit.data() + step * mBeamWidth * mVocabSize, - TRTDataType::value, ITensor::makeShape({mBeamWidth, mVocabSize})); - - for (int32_t bi = 0; bi < mBatchSize; ++bi) - { - auto logitsDeviceView = ITensor::slice(mDecodingWorkspace->getDeviceRuntimeLogits(), bi, 1); - mBufferManager->copy(*logitsHost, *logitsDeviceView); - } -} - -template -bool BaseSamplingLayerTest::checkResult(int32_t const* outputIds, std::vector> const& expectedIds) -{ - assert(expectedIds.size() == mMaxSeqLen * batchBeam()); - int failures = 0; - auto* const batchSlotsPtr = bufferCast(*mBatchSlots); - for (int32_t i = 0; i < mMaxSeqLen * mBatchSize; ++i) - { - int32_t s = i / mBatchSize; - int32_t b = i % mBatchSize; - auto const batchSlot = batchSlotsPtr[b]; - std::set expts = expectedIds.at(i); - auto const outputId = outputIds[batchSlot * mMaxSeqLen + s]; - if (expts.count(outputId) == 0) - { - if (failures < 10) - { - std::stringstream ss; - ss << " - Fail " - << " (step=" << s << ", batch=" << b << ") " - << "actual=" << outputId << ", expected"; - for (auto const& expt : expts) - { - ss << " " << expt; - } - TLLM_LOG_DEBUG("%s", ss.str().c_str()); - } - ++failures; - } - } - TLLM_LOG_DEBUG( - "check...%6s : failures: %d / %d", failures == 0 ? "....OK" : "FAILED", failures, mMaxSeqLen * batchBeam()); - return failures == 0; -} - -template -void BaseSamplingLayerTest::runTest( - std::vector> const& expectedOutputIds, TestSamplingParams const& params, int32_t endId) -{ - mBatchSize = params.batchSize; - if (params.beamWidth > 1) - { - mBeamWidth = params.beamWidth; - mMaxSeed = 1; - mComputeProbs = true; - } - initLayer(params); - - auto const decoderDomain - = tensorrt_llm::layers::DecoderDomain(maxBatchSize(), mBeamWidth, mVocabSize, mVocabSizePadded); - mDecodingWorkspace = std::make_unique( - mBufferManager, decoderDomain, TRTDataType::value, mSamplingLayer->getWorkspaceSize()); - mEndId = endId; - for (uint64_t seed = 0; seed < mMaxSeed; ++seed) - { - setup(seed, params); - - int32_t step = mMaxInputLen; - auto inputTensors = createInputTensors(step); - auto outputTensors = createOutputTensors(); - - for (step = mMaxInputLen; step < mMaxOutputLen; ++step) - { - // Reset by the test value since the sampling layer internally updates the logit buffer. - batchCopy(step); - if (params.isExternalDraftTokensLayerTest) - { - inputTensors = createInputTensors(step); - } - else - { - inputTensors->step = step; - } - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mSamplingLayer->forwardAsync(outputTensors, inputTensors, mDecodingWorkspace); - mStream->synchronize(); - } - - auto const outputIdsHost = mBufferManager->copyFrom(*mOutputIdsDevice, tensorrt_llm::runtime::MemoryType::kCPU); - - mStream->synchronize(); - - bool passed = checkResult(bufferCast(*outputIdsHost), expectedOutputIds); - EXPECT_TRUE(passed) << "Output ids check failed at seed " << seed; - if (!passed) - { - std::stringstream ss; - ss << "Actual output ids:" << std::endl << *outputIdsHost; - TLLM_LOG_DEBUG(ss.str()); - } - } -} - -template class BaseSamplingLayerTest; -template class BaseSamplingLayerTest; - -} // namespace tensorrt_llm::tests::layers::sampling diff --git a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h b/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h deleted file mode 100644 index 2bf782a5f5e8..000000000000 --- a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h +++ /dev/null @@ -1,191 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include - -#include - -#include "tensorrt_llm/layers/beamSearchLayer.h" -#include "tensorrt_llm/layers/externalDraftTokensLayer.h" -#include "tensorrt_llm/layers/samplingLayer.h" -#include "tensorrt_llm/layers/topKSamplingLayer.h" -#include "tensorrt_llm/layers/topPSamplingLayer.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -#include "tensorrt_llm/kernels/beamSearchKernels.h" -#include "tensorrt_llm/kernels/penaltyKernels.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/samplingTopPKernels.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -#include "tensorrt_llm/common/tllmException.h" - -namespace tensorrt_llm::tests::layers::sampling -{ - -constexpr float EPSILON = 1e-20f; - -template -void computeProb(T* probs, T const* logits, int batchSize, int vocabSize) -{ - // Compute the log probability from logits. - // logits = batchSize x vocabSize. - // probs = softmax(logits) (softmax along with vocab dimension) - // float is used for either T=float or half, since operations of half are - // not fully supported in a host function. - for (int bidx = 0; bidx < batchSize; ++bidx) - { - float maxval = -FLT_MAX; - for (int i = 0; i < vocabSize; ++i) - { - float logit = static_cast(logits[bidx * vocabSize + i]); - if (logit > maxval) - { - maxval = logit; - } - } - float sum = 0.0f; - for (int i = 0; i < vocabSize; ++i) - { - sum += expf(static_cast(logits[bidx * vocabSize + i]) - maxval); - } - for (int i = 0; i < vocabSize; ++i) - { - int idx = bidx * vocabSize + i; - float logit = static_cast(logits[idx]) - maxval; - probs[idx] = static_cast(expf(logit) / (sum + EPSILON)); - } - } -} - -struct TestSamplingParams -{ - std::vector topKs; - std::vector topPs; - std::vector temperatures; - std::vector repetitionPenalties; - std::vector presencePenalties; - std::vector frequencyPenalties; - std::vector minLengths; - std::vector decay; - std::vector minTopP; - std::vector topPResetIds; - int32_t batchSize = 6; - int32_t beamWidth = 1; - bool useBias = false; - bool isExternalDraftTokensLayerTest = false; - bool useDraftLogits = false; - bool isAirTopPExternalDraftTokensLayer = false; -}; - -template -class BaseSamplingLayerTest : public testing::Test -{ -protected: - using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; - using BufferPtr = tensorrt_llm::runtime::IBuffer::SharedPtr; - - static int32_t constexpr kDoubleBatchIdx = 2; - - int32_t seed = 0; - int32_t mBatchSize = -1; // setup by runTest - int32_t mBeamWidth = 1; - static int32_t constexpr mBatchSizeBadPad = 512; - uint64_t mMaxSeed = 32; - int32_t const mVocabSize = 8; - int32_t const mVocabSizePadded = mVocabSize; - - int32_t const mMaxInputLen = 0; // has no effect. - static int32_t constexpr mMaxOutputLen = 4; - int32_t const mMaxSeqLen = mMaxInputLen + mMaxOutputLen; - int32_t const mMaxTokensPerEngineStep = mMaxOutputLen; - - int32_t mEndId = mVocabSize; - - bool mComputeProbs = false; - - TensorPtr mContextLengthDevice; - TensorPtr mSeqLengthsDevice; - TensorPtr mFinishedDevice; - TensorPtr mOutputIdsDevice; - TensorPtr mEndIdsDevice; - TensorPtr mIdsPtrHost; - TensorPtr mBatchSlots; - - TensorPtr mEmbeddingBiasHost; - TensorPtr mEmbeddingBiasDevice; - - TensorPtr mCumLogProbsDevice; - TensorPtr mOutputLogProbsDevice; - - TensorPtr mCurandStatesDevice; - TensorPtr mPenaltyWorkspaceDevice; - - // For Beam Search - TensorPtr mSrcCacheIndirection; - TensorPtr mTgtCacheIndirection; - TensorPtr mParentIds; - TensorPtr mOutputIdsCBA; - TensorPtr mLogProbsCBA; - TensorPtr mSequenceLengthsCBA; - TensorPtr mCumLogProbsCBA; - TensorPtr mNormedScoresCBA; - TensorPtr mNumBeamsCBA; - TensorPtr mMinNormedScoresCBA; - TensorPtr mBatchDones; - TensorPtr mOutputIdsPtr; - TensorPtr mParentIdsPtr; - - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - std::shared_ptr mSamplingLayer; - std::shared_ptr mDecodingWorkspace; - - std::vector mTestLogitsInit; - - int32_t maxBatchSize() const - { - return kDoubleBatchIdx * mBatchSize; - } - - int32_t batchBeam() const - { - return mBatchSize * mBeamWidth; - } - - void setup(uint64_t seed, TestSamplingParams const& params); - - virtual void initLayer(TestSamplingParams const& params) = 0; - - virtual std::shared_ptr createInputTensors(int32_t step); - - std::shared_ptr createOutputTensors(); - - void batchCopy(int32_t step); - bool checkResult(int32_t const* outputIds, std::vector> const& expectedIds); - -public: - void runTest( - std::vector> const& expectedOutputIds, TestSamplingParams const& params, int32_t endId = -1); -}; - -typedef testing::Types FloatAndHalfTypes; - -} // namespace tensorrt_llm::tests::layers::sampling diff --git a/cpp/tests/unit_tests/layers/beamSearchLayerTest.cpp b/cpp/tests/unit_tests/layers/beamSearchLayerTest.cpp deleted file mode 100644 index 024a7cfdccf8..000000000000 --- a/cpp/tests/unit_tests/layers/beamSearchLayerTest.cpp +++ /dev/null @@ -1,70 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/executor/types.h" -#include "tests/unit_tests/layers/baseSamplingLayerTest.h" - -namespace -{ - -namespace tle = tensorrt_llm::executor; - -using namespace tensorrt_llm::tests::layers::sampling; -using namespace tensorrt_llm::runtime; - -template -class BeamSearchLayerTest : public BaseSamplingLayerTest -{ - void SetUp() override - { - this->mStream = std::make_shared(); - this->mBufferManager = std::make_shared(this->mStream); - } - - void initLayer(TestSamplingParams const& params) override - { - auto decodingMode = tle::DecodingMode::BeamSearch(); - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain( - this->maxBatchSize(), params.beamWidth, this->mVocabSize, this->mVocabSizePadded); - this->mSamplingLayer = std::make_shared>( - decodingMode, decodingDomain, this->mBufferManager); - } -}; - -TYPED_TEST_SUITE(BeamSearchLayerTest, FloatAndHalfTypes); - -TYPED_TEST(BeamSearchLayerTest, BeamWidth2) -{ - TestSamplingParams params; - - params.batchSize = 3; - params.beamWidth = 2; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, // step 3 - {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -} // namespace diff --git a/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp b/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp deleted file mode 100644 index cc3b9c411f5b..000000000000 --- a/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp +++ /dev/null @@ -1,2197 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tests/unit_tests/layers/dynamicDecodeLayerTest.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include - -namespace tensorrt_llm::tests::layers::sampling -{ - -// TODO: -// Add tests for -// - finished states -// - finished sum -// - max length -// - padded vocab -// - beam search - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using namespace tensorrt_llm::common; - -namespace tk = tensorrt_llm::kernels; -namespace trk = tensorrt_llm::runtime::kernels; -namespace tle = tensorrt_llm::executor; - -constexpr float EPSILON = 1e-20f; - -inline bool almostEqual(float a, float b, float atol = 1e-5, float rtol = 1e-8) -{ - // Params: a = value to compare and b = reference - // This function follows implementation of numpy.isclose(), which checks - // abs(a - b) <= (atol + rtol * abs(b)). - // Note that the inequality above is asymmetric where b is considered as - // a reference value. To account into both absolute/relative errors, it - // uses absolute tolerance and relative tolerance at the same time. The - // default values of atol and rtol borrowed from numpy.isclose(). For the - // case of nan value, the result will be true. - if (isnan(a) && isnan(b)) - { - return true; - } - - if (isinf(a) && isinf(b)) - { - return true; - } - return fabs(a - b) <= (atol + rtol * fabs(b)); -} - -template -bool compareValues(T* out, T* ref, size_t size) -{ - bool isFp32 = sizeof(T) == 4; - float atol = isFp32 ? 1e-4f : 1e-3f; - float rtol = isFp32 ? 1e-2f : 1e-1f; - - size_t failures = 0; - float relativeGap = 0.0f; - - for (size_t i = 0; i < size; ++i) - { - // The values for the output and the reference. - float a = (float) out[i]; - float b = (float) ref[i]; - - bool ok = almostEqual(a, b, atol, rtol); - // Print the error. - if (!ok && failures < 4) - { - TLLM_LOG_DEBUG(">> invalid result for i=%lu:", i); - TLLM_LOG_DEBUG(">> found......: %10.6f", a); - TLLM_LOG_DEBUG(">> expected...: %10.6f", b); - TLLM_LOG_DEBUG(">> error......: %.6f", fabsf(a - b)); - TLLM_LOG_DEBUG(">> tol........: %.6f", atol + rtol * fabs(b)); - } - // Update the number of failures. - failures += ok ? 0 : 1; - // Update the relative gap. - relativeGap += fabsf(a - b) / (fabsf(b) + EPSILON); - } - - relativeGap /= size; - - // Allow not matched up to 0% elements. - size_t tolFailures = (size_t) (0.0 * size); - TLLM_LOG_DEBUG("check... : %-50s (failures: %.2f%% atol: %.2e rtol: %.2e rel_gap: %.2e%%)", - failures <= tolFailures ? "....OK" : "FAILED", 100. * failures / size, atol, rtol, 100. * relativeGap); - return failures <= tolFailures; -} - -template bool compareValues(float* out, float* ref, size_t size); -template bool compareValues(half* out, half* ref, size_t size); - -template -void DynamicDecodeLayerTest::SetUp() -{ - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); -} - -template -void DynamicDecodeLayerTest::allocateData(TestSamplingParams const& params, TokenIdType endId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mEndId = endId == -1 ? mVocabSize - 1 : endId; - - mDecodingMode = params.decodingMode.value_or( - [this]() - { - if (this->mBeamWidth == 1) - { - return tle::DecodingMode::TopKTopP(); - } - else - { - return tle::DecodingMode::BeamSearch(); - } - }()); - - mMaxTokensPerStep = mDecodingMode.isMedusa() ? mMaxOutputLen - mMaxInputLen : 1; - - auto speculativeDecodingModule = std::make_shared( - params.maxNumMedusaHeads.value_or(0), mMaxTokensPerStep - 1, mMaxTokensPerStep); - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain( - mMaxBatchSize, mBeamWidth, mVocabSize, mVocabSizePadded, speculativeDecodingModule); - - mDecodeLayer - = std::make_unique>(mDecodingMode, decodingDomain, mBufferManager); - - auto const dataType = TRTDataType::value; - - mLogitsDevice = mBufferManager->gpu( - ITensor::makeShape({mBatchSize, mMaxTokensPerStep, mBeamWidth, mVocabSizePadded}), dataType); - mRuntimeLogitsHost - = BufferManager::pinned(ITensor::makeShape({mBatchSize, mBeamWidth, mVocabSizePadded}), dataType); - - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mFinishedDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize}), TRTDataType::value); - mFinishedSumDevice = BufferManager::pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); - mOutputIdsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mNewTokens - = BufferManager::pinned(ITensor::makeShape({mMaxTokensPerStep, mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mEmbeddingBiasHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mVocabSizePadded}), dataType); - mEmbeddingBiasDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mVocabSizePadded}), dataType); - - mRefLogProbsHost - = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kFLOAT); - mOutputLogProbsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kFLOAT); - mOutputLogProbsTiledDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxSeqLen, mMaxBatchSize}), tensorrt_llm::DataType::kFLOAT); - - mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kFLOAT); - - mMaxBadWordsLen = getMaxWordsLen(params.badWords); - mMaxStopWordsLen = getMaxWordsLen(params.stopWords); - - mBadWords = BufferManager::pinned( - ITensor::makeShape({mMaxBatchSize, 2, mMaxBadWordsLen}), tensorrt_llm::DataType::kINT32); - mBadWordsLens = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mBadWordsPtrs = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT64); - - mStopWords = BufferManager::pinned( - ITensor::makeShape({mMaxBatchSize, 2, mMaxStopWordsLen}), tensorrt_llm::DataType::kINT32); - mStopWordsLens = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mStopWordsPtrs = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT64); - - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); - - if (mDecodingMode.isMedusa()) - { - allocateMedusaData(params); - } - mDecodingWorkspace = std::make_unique( - mBufferManager, decodingDomain, TRTDataType::value, mDecodeLayer->getWorkspaceSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayerTest::allocateMedusaData(TestSamplingParams const& params) -{ - auto const dataType = TRTDataType::value; - mMaxMedusaHeads = params.maxNumMedusaHeads.value(); - mPathsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep, mMaxMedusaHeads + 1}), tensorrt_llm::DataType::kINT32); - mAcceptedLengths = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mMedusaLogitsDevice = BufferManager::pinned( - ITensor::makeShape({mMaxMedusaHeads, mMaxBatchSize, mMaxTokensPerStep, mVocabSizePadded}), dataType); - mNextDraftTokensDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep - 1}), tensorrt_llm::DataType::kINT32); - mTokensPerStepDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mTreeIdsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep - 1}), tensorrt_llm::DataType::kINT32); - mAcceptedLengthCumSumDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize + 1}), tensorrt_llm::DataType::kINT32); - mPackedPathsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize * mMaxMedusaHeads}), tensorrt_llm::DataType::kINT32); -} - -template -void DynamicDecodeLayerTest::setup(uint64_t seed, TestSamplingParams const& params) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const dataType = TRTDataType::value; - - // clang-format off - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1, 0.0) - mTestLogitsInit = { - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // step 0 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // step 1 - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // step 2 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX // step 3 - }; - - // clang-format on - - trk::invokeFill(*mSeqLengthsDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mContextLengthDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mFinishedDevice, uint8_t{0}, *mStream); - trk::invokeFill(*mOutputIdsDevice, TokenIdType{0}, *mStream); - trk::invokeFill(*mEmbeddingBiasDevice, T{0.0f}, *mStream); - trk::invokeFill(*mCumLogProbsDevice, float{0.0f}, *mStream); - trk::invokeFill(*mOutputLogProbsDevice, float{0.0f}, *mStream); - trk::invokeFill(*mOutputLogProbsTiledDevice, float{0.0f}, *mStream); - trk::invokeFill(*mRefLogProbsHost, float{0.0f}, *mStream); - trk::invokeFill(*mEndIdsDevice, TokenIdType{mEndId}, *mStream); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - batchSlotsPtr[bi] = 2 * bi; - } - - if (params.useBias) - { - auto embeddingBiasHostPtr = bufferCast(*mEmbeddingBiasHost); - for (SizeType32 bi = 0; bi < mMaxBatchSize; bi++) - { - for (SizeType32 vi = 0; vi < mVocabSizePadded; vi++) - { - embeddingBiasHostPtr[bi * mVocabSizePadded + vi] = 2 <= vi && vi < 6 ? T{2.0f} : T{0.0f}; - } - } - mBufferManager->copy(*mEmbeddingBiasHost, *mEmbeddingBiasDevice); - } - - mLogitsVec.resize(mBatchSize); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - mLogitsVec[bi] = ITensor::slice(mLogitsDevice, bi, 1); - } - - if (mDecodingMode.isMedusa()) - { - auto const maxMedusaHeads = params.maxNumMedusaHeads.value(); - - trk::invokeFill(*mPathsDevice, SizeType32{-1}, *mStream); - trk::invokeFill(*mAcceptedLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mNextDraftTokensDevice, TokenIdType{mEndId}, *mStream); - trk::invokeFill(*mTokensPerStepDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mTreeIdsDevice, SizeType32{0}, *mStream); - - auto const logitsHost - = ITensor::wrap(mTestLogitsInit, ITensor::makeShape({mMaxTokensPerStep, mVocabSizePadded})); - for (SizeType32 hi = 0; hi < maxMedusaHeads; ++hi) - { - TensorPtr logitsHeadDeviceView = ITensor::slice(mMedusaLogitsDevice, hi, 1); - logitsHeadDeviceView->squeeze(0); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - TensorPtr logitsHeadBatchDeviceView = ITensor::slice(logitsHeadDeviceView, bi, 1); - mBufferManager->copy(*logitsHost, *logitsHeadBatchDeviceView); - } - } - - auto paths = params.paths.value(); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const numPaths = static_cast(paths[bi].size() / (maxMedusaHeads + 1)); - auto const pathsHost = ITensor::wrap(paths[bi], ITensor::makeShape({1, numPaths, maxMedusaHeads + 1})); - TensorPtr pathsDeviceSlice = ITensor::slice(mPathsDevice, batchSlotsPtr[bi], 1); - pathsDeviceSlice->squeeze(0); - TensorPtr pathsNumPathsDeviceSlice = ITensor::slice(pathsDeviceSlice, 0, numPaths); - pathsNumPathsDeviceSlice->unsqueeze(0); - mBufferManager->copy(*pathsHost, *pathsNumPathsDeviceSlice); - } - - auto tokensPerStep = params.tokensPerStep.value(); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - TensorPtr tokensPerStepDeviceSlice = ITensor::slice(mTokensPerStepDevice, batchSlotsPtr[bi], 1); - trk::invokeFill(*tokensPerStepDeviceSlice, SizeType32{tokensPerStep[bi]}, *mStream); - } - - auto outputIds = params.outputIds.value(); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const outputIdsBatchHost = ITensor::wrap(outputIds[bi], ITensor::makeShape({mMaxTokensPerStep - 1})); - - auto outputIdsDevice = ITensor::slice(mNextDraftTokensDevice, batchSlotsPtr[bi], 1); - mBufferManager->copy(*outputIdsBatchHost, *outputIdsDevice); - } - } - - auto setupParams = std::make_shared(); - setupParams->penaltyParams = std::make_shared(); - setupParams->penaltyParams->temperature - = params.temperatures.size() ? std::make_optional>(params.temperatures) : std::nullopt; - setupParams->penaltyParams->repetitionPenalty = params.repetitionPenalties.size() - ? std::make_optional>(params.repetitionPenalties) - : std::nullopt; - setupParams->penaltyParams->presencePenalty = params.presencePenalties.size() - ? std::make_optional>(params.presencePenalties) - : std::nullopt; - setupParams->penaltyParams->frequencyPenalty = params.frequencyPenalties.size() - ? std::make_optional>(params.frequencyPenalties) - : std::nullopt; - setupParams->penaltyParams->minLength - = params.minLengths.size() ? std::make_optional>(params.minLengths) : std::nullopt; - - setupParams->banWordsParams = std::make_shared(); - setupParams->banWordsParams->noRepeatNgramSize = params.repeatNGramSizes.size() - ? std::make_optional>(params.repeatNGramSizes) - : std::nullopt; - - if (mDecodingMode.isTopKorTopP()) - { - auto samplingParams = std::make_shared(); - samplingParams->randomSeed = std::make_optional>({seed}); - samplingParams->runtimeTopK - = params.topKs.size() ? std::make_optional>(params.topKs) : std::nullopt; - samplingParams->runtimeTopP - = params.topPs.size() ? std::make_optional>(params.topPs) : std::nullopt; - samplingParams->topPDecay - = params.decay.size() ? std::make_optional>(params.decay) : std::nullopt; - samplingParams->topPMin - = params.minTopP.size() ? std::make_optional>(params.minTopP) : std::nullopt; - samplingParams->topPResetIds = params.topPResetIds.size() - ? std::make_optional>(params.topPResetIds) - : std::nullopt; - samplingParams->normalizeLogProbs = {false}; - samplingParams->outputLogProbs = {true}; - samplingParams->cumLogProbs = {true}; - - setupParams->decodingParams = samplingParams; - } - else if (mDecodingMode.isMedusa()) - { - auto medusaParams = std::make_shared(); - medusaParams->runtimeHeadsTopK = params.topKMedusaHeads; - medusaParams->randomSeed = std::make_optional>({seed}); - medusaParams->runtimeTopK - = params.topKs.size() ? std::make_optional>(params.topKs) : std::nullopt; - - setupParams->decodingParams = medusaParams; - } - - initXWordsTensors(batchSlotsPtr, bufferCast(*mBadWords), - reinterpret_cast(bufferCast(*mBadWordsPtrs)), bufferCast(*mBadWordsLens), - mMaxBadWordsLen, params.badWords); - initXWordsTensors(batchSlotsPtr, bufferCast(*mStopWords), - reinterpret_cast(bufferCast(*mStopWordsPtrs)), bufferCast(*mStopWordsLens), - mMaxStopWordsLen, params.stopWords); - mDecodeLayer->setup(mBatchSize, mBeamWidth, mBatchSlots, setupParams, mDecodingWorkspace); - - mStream->synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -SizeType32 DynamicDecodeLayerTest::getMaxWordsLen( - std::vector>> const& inputWords) -{ - SizeType32 maxWordsLen = 0; - for (auto const& batchWords : inputWords) - { - SizeType32 wordsLen = 0; - for (auto const& words : batchWords) - { - wordsLen += words.size(); - } - if (wordsLen == batchWords.size()) - { - wordsLen += 1; - } - maxWordsLen = std::max(maxWordsLen, wordsLen); - } - return maxWordsLen; -} - -template -void DynamicDecodeLayerTest::initXWordsTensors(SizeType32* batchSlotsPtr, SizeType32* wordsData, - SizeType32** wordsPtr, SizeType32* wordsLenData, SizeType32 maxWordsLen, - std::vector>> const& inputWords) -{ - std::fill(wordsData, wordsData + mMaxBatchSize * 2 * maxWordsLen, -1); - for (SizeType32 bi = 0; bi < inputWords.size(); bi++) - { - auto const batchSlot = batchSlotsPtr[bi]; - SizeType32 totalLen = 0; - for (SizeType32 wi = 0; wi < inputWords[bi].size(); ++wi) - { - for (SizeType32 si = 0; si < inputWords[bi][wi].size(); ++si) - { - wordsData[batchSlot * 2 * maxWordsLen + 0 * maxWordsLen + totalLen + si] = inputWords[bi][wi][si]; - } - totalLen += inputWords[bi][wi].size(); - // Do not add value if words is empty - if (totalLen > 0) - { - wordsData[batchSlot * 2 * maxWordsLen + 1 * maxWordsLen + wi] = totalLen; - } - } - } - - for (SizeType32 bi = 0; bi < inputWords.size(); bi++) - { - auto const batchSlot = batchSlotsPtr[bi]; - wordsPtr[batchSlot] = wordsData + batchSlot * 2 * maxWordsLen; - - wordsLenData[batchSlot] = maxWordsLen; - } -} - -template -void DynamicDecodeLayerTest::createMedusaInputs(std::shared_ptr& baseInputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto inputs = std::dynamic_pointer_cast(baseInputs); - - auto batchSlots = BufferRange(*mBatchSlots); - std::vector> medusaLogits(mMaxBatchSize); - auto const medusaLogitsPtr = bufferCast(*mMedusaLogitsDevice); - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - medusaLogits[bi].resize(mMaxMedusaHeads); - } - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - for (SizeType32 hi = 0; hi < mMaxMedusaHeads; ++hi) - { - TensorPtr logitsHead = ITensor::slice(mMedusaLogitsDevice, hi, 1); - logitsHead->squeeze(0); - TensorPtr logitsHeadBatch = ITensor::slice(logitsHead, bi, 1); - medusaLogits[batchSlots[bi]][hi] = logitsHeadBatch; - } - } - - inputs->paths = mPathsDevice; - inputs->treeIds = mTreeIdsDevice; - inputs->medusaLogits = medusaLogits; - inputs->curTokensPerStep = mTokensPerStepDevice; - inputs->targetTokensPerStep = mTokensPerStepDevice; - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -std::shared_ptr DynamicDecodeLayerTest::createInputTensors(SizeType32 step) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 constexpr ite = 0; - std::shared_ptr forwardParams; - if (mDecodingMode.isTopKorTopP()) - { - forwardParams = std::make_shared(mEndIdsDevice, mBatchSlots, step, ite, mBatchSize); - } - else if (mDecodingMode.isMedusa()) - { - forwardParams = std::make_shared(mEndIdsDevice, mBatchSlots, mBatchSize); - } - - forwardParams->embeddingBias = mEmbeddingBiasDevice; - - forwardParams->finished = mFinishedDevice; - - if (mUseLogitsVec) - { - forwardParams->logitsVec = mLogitsVec; - } - else - { - forwardParams->logits = mLogitsDevice; - } - - forwardParams->banWordsInputs = std::make_shared(mBatchSize); - forwardParams->banWordsInputs->badWordsPtr = mBadWordsPtrs; - forwardParams->banWordsInputs->badWordsLengths = mBadWordsLens; - forwardParams->banWordsInputs->maxBadWordsLen = mMaxBadWordsLen; - - forwardParams->stopCriteriaInputs = std::make_shared(mBatchSize); - forwardParams->stopCriteriaInputs->stopWordsPtr = mStopWordsPtrs; - forwardParams->stopCriteriaInputs->stopWordsLengths = mStopWordsLens; - forwardParams->stopCriteriaInputs->maxStopWordsLen = mMaxStopWordsLen; - - if (mDecodingMode.isMedusa()) - { - createMedusaInputs(forwardParams); - } - - // TODO: extend to - // std::optional src_cache_indirection; - // std::optional sequence_limit_length; - // std::optional input_lengths; - // std::optional> logitsVec; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return forwardParams; -} - -template -void DynamicDecodeLayerTest::createMedusaOutputs(std::shared_ptr& baseOutputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto outputs = std::dynamic_pointer_cast(baseOutputs); - outputs->nextDraftTokens = mNextDraftTokensDevice; - - outputs->numNewTokens = mAcceptedLengths; - - outputs->numNewTokensCumSum = mAcceptedLengthCumSumDevice; - - outputs->pathsOffsets = mPackedPathsDevice; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -std::shared_ptr DynamicDecodeLayerTest::createOutputTensors() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - std::shared_ptr outputParams; - - if (mDecodingMode.isMedusa()) - { - outputParams = std::make_shared(mOutputIdsDevice); - } - else - { - outputParams = std::make_shared(mOutputIdsDevice); - } - - outputParams->sequenceLength = mSeqLengthsDevice; - - outputParams->finished = mFinishedDevice; - - outputParams->finishedSum = mFinishedSumDevice; - - outputParams->newTokens = mNewTokens; - - if (!mDecodingMode.isMedusa()) - { - // Output log probs are not supported in Medusa - outputParams->cumLogProbs = mCumLogProbsDevice; - - outputParams->outputLogProbs = mOutputLogProbsDevice; - - outputParams->outputLogProbsTiled = mOutputLogProbsTiledDevice; - } - - if (mDecodingMode.isMedusa()) - { - createMedusaOutputs(outputParams); - } - - // TODO: extend to - // std::optional parent_ids; - // std::optional tgt_cache_indirection; - // std::shared_ptr beamHypotheses; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return outputParams; -} - -template -void DynamicDecodeLayerTest::batchCopy(SizeType32 step) -{ - auto const logitsHost = ITensor::wrap(mTestLogitsInit.data() + step * mVocabSizePadded, - std::is_same_v ? tensorrt_llm::DataType::kFLOAT : tensorrt_llm::DataType::kHALF, - ITensor::makeShape({mMaxTokensPerStep, mVocabSizePadded})); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - TensorPtr logitsDeviceView = ITensor::slice(mLogitsDevice, bi, 1); - logitsDeviceView->squeeze(0); - mBufferManager->copy(*logitsHost, *logitsDeviceView); - } - mLogitsRefHost = mBufferManager->copyFrom(*mLogitsDevice, tensorrt_llm::runtime::MemoryType::kCPU); -} - -template -bool DynamicDecodeLayerTest::checkResult(TokenIdType* outputIds, - std::vector> const& expectedIds, SizeType32* seqLens, SizeType32 leadingDim, - SizeType32 stride, SizeType32 step, bool outputIdsTransposed, SizeType32 strideTransposed) -{ - SizeType32 failures = 0; - auto const batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 i = 0; i < leadingDim * stride; ++i) - { - auto const s = i / stride; - auto const b = i % stride; - auto const batchSlot = batchSlotsPtr[b]; - if (seqLens[batchSlot] <= step + s) - { - continue; - } - auto const& expts = expectedIds.at(i + step * stride); - auto const outputIdIdx = outputIdsTransposed ? s * strideTransposed + batchSlot : batchSlot * leadingDim + s; - auto const outputId = outputIds[outputIdIdx]; - if (expts.count(outputId) == 0) - { - if (failures < 10) - { - std::stringstream ss; - ss << " - Fail " - << " (step=" << s << ", batch=" << b << ") " - << "actual=" << outputId << ", expected"; - for (auto const& expt : expts) - { - ss << " " << expt; - } - TLLM_LOG_DEBUG("%s", ss.str().c_str()); - } - ++failures; - } - } - TLLM_LOG_DEBUG( - "check...%6s : failures: %d / %d", failures == 0 ? "....OK" : "FAILED", failures, leadingDim * stride); - return failures == 0; -} - -template -void DynamicDecodeLayerTest::fillRefLogits( - SizeType32 const* seqLenHost, std::vector> const& expectedOutputIds, SizeType32 step) -{ - auto const batchSlotsPtr = bufferCast(*mBatchSlots); - auto const runtimeLogitsHost = bufferCast(*mRuntimeLogitsHost); - for (SizeType32 bi = 0; bi < mBatchBeam; ++bi) - { - auto const batchSlot = batchSlotsPtr[bi]; - if (seqLenHost[batchSlot] <= step) - { - continue; - } - auto& expectedSet = expectedOutputIds[step * mBatchBeam + bi]; - TLLM_CHECK(expectedSet.size() == 1); - auto expectedToken = *expectedSet.begin(); - bufferCast(*mRefLogProbsHost)[batchSlot * mMaxSeqLen + step] - = logf(runtimeLogitsHost[bi * mVocabSizePadded + expectedToken]); - } -} - -template -void DynamicDecodeLayerTest::runTestImpl( - std::vector> const& expectedOutputIds, TestSamplingParams const& params, TokenIdType endId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - bool greedySearch - = std::all_of(expectedOutputIds.begin(), expectedOutputIds.end(), [](auto v) { return v.size() == 1; }); - for (uint64_t seed = 0; seed < mMaxSeed; ++seed) - { - setup(seed, params); - - auto step = mMaxInputLen; - auto inputTensors = createInputTensors(step); - auto outputTensors = createOutputTensors(); - - for (step = mMaxInputLen; step < mMaxOutputLen; step += mMaxTokensPerStep) - { - // Reset by the test value since the sampling layer internally update the logit buffer. - batchCopy(step); - if (mUseLogitsVec) - { - inputTensors->logitsVec = mLogitsVec; - inputTensors->logits = std::nullopt; - } - else - { - inputTensors->logits = mLogitsDevice; - inputTensors->logitsVec = std::nullopt; - } - inputTensors->step = step; - mDecodeLayer->forwardAsync(outputTensors, inputTensors, mDecodingWorkspace); - mStream->synchronize(); - auto const newTokensHost = mBufferManager->copyFrom(*mNewTokens, tensorrt_llm::runtime::MemoryType::kCPU); - auto const seqLenHost - = mBufferManager->copyFrom(*mSeqLengthsDevice, tensorrt_llm::runtime::MemoryType::kCPU); - auto const logitsHost = mBufferManager->copyFrom(*mLogitsDevice, tensorrt_llm::runtime::MemoryType::kCPU); - mBufferManager->copy(mDecodingWorkspace->getDeviceRuntimeLogits()->data(), *mRuntimeLogitsHost, - tensorrt_llm::runtime::MemoryType::kGPU); - mStream->synchronize(); - - if (greedySearch && !mDecodingMode.isMedusa()) - { - fillRefLogits(bufferCast(*seqLenHost), expectedOutputIds, step); - } - - { - auto const passed = checkResult(bufferCast(*newTokensHost), expectedOutputIds, - bufferCast(*seqLenHost), mMaxTokensPerStep, mBatchBeam, step, /* transposed */ true, - /* stride transposed */ mMaxBatchSize * mBeamWidth); - EXPECT_TRUE(passed) << "New tokens check failed at seed " << seed; - if (!passed) - { - std::stringstream ss; - ss << "New tokens ids:" << std::endl << *newTokensHost; - TLLM_LOG_DEBUG(ss.str()); - } - } - - // Check if logits were not modified in-place - { - auto const passed = compareValues(bufferCast(*mLogitsRefHost), bufferCast(*logitsHost), - mBatchSize * mMaxTokensPerStep * mBeamWidth * mVocabSizePadded); - EXPECT_TRUE(passed) << "Unmodified logits check failed at seed " << seed; - } - } - - auto const outputIdsHost = mBufferManager->copyFrom(*mOutputIdsDevice, tensorrt_llm::runtime::MemoryType::kCPU); - auto const seqLenHost = mBufferManager->copyFrom(*mSeqLengthsDevice, tensorrt_llm::runtime::MemoryType::kCPU); - auto const logProbsHost - = mBufferManager->copyFrom(*mOutputLogProbsDevice, tensorrt_llm::runtime::MemoryType::kCPU); - - mStream->synchronize(); - - { - auto const passed = checkResult(bufferCast(*outputIdsHost), expectedOutputIds, - bufferCast(*seqLenHost), mMaxSeqLen, mBatchBeam, /* step */ 0); - EXPECT_TRUE(passed) << "Output Ids check failed at seed " << seed; - if (!passed) - { - std::stringstream ss; - ss << "Actual output ids:" << std::endl << *outputIdsHost; - TLLM_LOG_DEBUG(ss.str()); - } - } - - if (greedySearch && !mDecodingMode.isMedusa()) - { - auto const passed = compareValues( - bufferCast(*logProbsHost), bufferCast(*mRefLogProbsHost), mMaxSeqLen * mMaxBatchSize); - EXPECT_TRUE(passed) << "Log probs check failed at seed " << seed; - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void DynamicDecodeLayerTest::runTest( - std::vector> const& expectedOutputIds, TestSamplingParams const& params, TokenIdType endId) -{ - allocateData(params, endId); - - if (!params.decodingMode.has_value() || !params.decodingMode->isMedusa()) - { - TLLM_LOG_DEBUG("Run test with linear logits"); - mUseLogitsVec = false; - runTestImpl(expectedOutputIds, params, endId); - } - TLLM_LOG_DEBUG("Run test with vectorized logits"); - mUseLogitsVec = true; - runTestImpl(expectedOutputIds, params, endId); -} - -template class DynamicDecodeLayerTest; -template class DynamicDecodeLayerTest; - -TYPED_TEST_SUITE(DynamicDecodeLayerTest, FloatAndHalfTypes); - -TYPED_TEST(DynamicDecodeLayerTest, TopK) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopK1TopP0) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, BatchTopK) -{ - std::vector topKs = {2, 1, 1, 2, 1, 1}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {2, 3}, {2}, {2}, {2, 3}, {2}, {2}, // step 2 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKTopP) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, BatchTopKTopP) -{ - std::vector topKs = {2, 2, 1, 2, 2, 1}; - float topP = 0.3; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKBatchTopP) -{ - SizeType32 topK = 2; - std::vector topPs = {0.5, 0.3, 0.5, 0.5, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4, 5}, {4, 5}, {4}, {4, 5}, // step 0 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0, 1}, // step 1 - {2, 3}, {2}, {2, 3}, {2, 3}, {2}, {2, 3}, // step 2 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, BatchTopKBatchTopP) -{ - std::vector topKs = {2, 2, 0, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0}, // step 1 - {2, 3}, {2}, {2, 3}, {2, 3}, {2}, {2}, // step 2 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsZeroTopK) -{ - SizeType32 topK = 0; - TestSamplingParams params; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsZeroTopP) -{ - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsZeroTopKTopP) -{ - SizeType32 topK = 0; - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsZeroBatchTopKTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsZeroTopKBatchTopP) -{ - SizeType32 topK = 0; - std::vector topPs = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - TestSamplingParams params; - params.topPs = topPs; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsBatchTopKContainZero) -{ - std::vector topKs = {2, 1, 0, 0, 2, 1}; - TestSamplingParams params; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsBatchTopKTopPContainZero) -{ - std::vector topKs = {2, 2, 1, 0, 2, 0}; - float topP = 0.0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2, 3}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, InvalidArgsBatchTopKBatchTopPContainZero) -{ - std::vector topKs = {0, 2, 1, 2, 2, 0}; - std::vector topPs = {0.0, 0.3, 0.9, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topPs = topPs; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4, 5}, {4}, {4, 5}, // step 0 - {0}, {0}, {0}, {0, 1}, {0}, {0, 1}, // step 1 - {2}, {2}, {2}, {2, 3}, {2}, {2, 3}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPTemperature) -{ - float temperature = 0.01f; - TestSamplingParams params; - params.temperatures = {temperature}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPTemperatureNoTemperatureMode) -{ - float temperature = 0.01f; - TestSamplingParams params; - params.temperatures = {temperature}; - params.topPs = {1.0f}; - params.decodingMode = tle::DecodingMode::TopP().useTemperature(false); - std::vector> expectedOutputIds{ - {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5, 6, 7}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, // step 1 - {2, 3, 4, 5}, {2, 3, 4, 5}, {2, 3, 4, 5}, {2, 3, 4, 5}, {2, 3, 4, 5}, {2, 3, 4, 5}, // step 2 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPTemperatureBatch) -{ - std::vector temperatures = {0.01f, 1e3f, 1.0f, 1.0f, 0.01f, 1.0f}; - TestSamplingParams params; - params.temperatures = temperatures; - params.topPs = {0.5f}; - std::vector> expectedOutputIds{ - {4}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4, 5}, // step 0 - {0}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0, 1}, // step 1 - {2}, {2, 3, 4, 5}, {2, 3}, {2, 3}, {2}, {2, 3}, // step 2 - {0}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPTemperatureMultipleRequests) -{ - this->allocateData(TestSamplingParams{}); - { - std::vector temperatures = {0.01f, 1e3f, 1.0f, 1.0f, 0.01f, 1.0f}; - TestSamplingParams params; - params.temperatures = temperatures; - params.topPs = {0.5f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4, 5}, // step 0 - {0}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0, 1}, // step 1 - {2}, {2, 3, 4, 5}, {2, 3}, {2, 3}, {2}, {2, 3}, // step 2 - {0}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - TestSamplingParams params; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - float temperature = 1.0f; - TestSamplingParams params; - params.temperatures = {temperature}; - params.topPs = {0.5f}; - std::vector> expectedOutputIds{ - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionPenalty) -{ - SizeType32 topK = 1; - float repetitionPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionPenaltyNoRepetitionMode) -{ - SizeType32 topK = 1; - float repetitionPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.topPs = {0.3f}; - params.decodingMode = tle::DecodingMode::TopP().useOccurrencePenalties(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionPenaltiesBatch) -{ - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = repetitionPenalties; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionPenaltyMultipleRequests) -{ - this->allocateData(TestSamplingParams{}); - { - float repetitionPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - TestSamplingParams params; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = repetitionPenalties; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPPresencePenalty) -{ - float presencePenalty = 1e9f; - TestSamplingParams params; - params.presencePenalties = {presencePenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPPresencePenaltyNoPresenceMode) -{ - float presencePenalty = 1e9f; - TestSamplingParams params; - params.presencePenalties = {presencePenalty}; - params.topPs = {0.3f}; - params.decodingMode = tle::DecodingMode::TopP().useOccurrencePenalties(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPPresencePenaltiesBatch) -{ - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.presencePenalties = presencePenalties; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPPresencePenaltyMultipleRequests) -{ - this->allocateData(TestSamplingParams{}); - { - float presencePenalty = 1e9f; - TestSamplingParams params; - params.presencePenalties = {presencePenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - TestSamplingParams params; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.presencePenalties = presencePenalties; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPFrequencyPenalty) -{ - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.frequencyPenalties = {frequencyPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPFrequencyPenaltyNoFrequencyMode) -{ - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.frequencyPenalties = {frequencyPenalty}; - params.topPs = {0.3f}; - params.decodingMode = tle::DecodingMode::TopP().useOccurrencePenalties(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPFrequencyPenaltiesBatch) -{ - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.frequencyPenalties = frequencyPenalties; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPFrequencyPenaltyMultipleRequests) -{ - this->allocateData(TestSamplingParams{}); - { - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.frequencyPenalties = {frequencyPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - TestSamplingParams params; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.frequencyPenalties = frequencyPenalties; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionPresencePenalty) -{ - float repetitionPenalty = 1e9f; - float presencePenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.presencePenalties = {presencePenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionPresencePenaltiesBatch) -{ - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalties}; - params.presencePenalties = {presencePenalties}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionFrequencyPenalty) -{ - float repetitionPenalty = 1e9f; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.frequencyPenalties = {frequencyPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPRepetitionFrequencyPenaltiesBatch) -{ - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalties}; - params.frequencyPenalties = {frequencyPenalties}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPPresenceFrequencyPenalty) -{ - float presencePenalty = 1e9f; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.presencePenalties = {presencePenalty}; - params.frequencyPenalties = {frequencyPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPPresenceFrequencyPenaltiesBatch) -{ - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.presencePenalties = {presencePenalties}; - params.frequencyPenalties = {frequencyPenalties}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPFullPenalty) -{ - float repetitionPenalty = 1e9f; - float presencePenalty = 1e9f; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.presencePenalties = {presencePenalty}; - params.frequencyPenalties = {frequencyPenalty}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPFullPenaltiesBatch) -{ - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalties}; - params.presencePenalties = {presencePenalties}; - params.frequencyPenalties = {frequencyPenalties}; - params.topPs = {0.3f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPMinLengthBatch) -{ - std::vector minLengths = {3, 1, 1, 3, 0, 3}; - TestSamplingParams params; - params.minLengths = minLengths; - params.topPs = {0.3f}; - TokenIdType const endId = 0; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {1}, {0}, {0}, {1}, {0}, {1}, // step 1 - {2}, {0}, {0}, {2}, {0}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params, endId); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPMinLengthBatchNoMinLengthMode) -{ - std::vector minLengths = {3, 1, 1, 3, 0, 3}; - TestSamplingParams params; - params.minLengths = minLengths; - params.topPs = {0.3f}; - TokenIdType const endId = 0; - params.decodingMode = tle::DecodingMode::TopP().useMinLength(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params, endId); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopPBias) -{ - TestSamplingParams params; - params.topPs = {0.5f}; - params.useBias = true; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKTemperature) -{ - SizeType32 topK = 2; - float temperature = 0.01f; - TestSamplingParams params; - params.temperatures = {temperature}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKTemperatureBatch) -{ - SizeType32 topK = 2; - std::vector temperatures = {0.01f, 1e3f, 1.0f, 0.5f, 0.01f, 1.0f}; - TestSamplingParams params; - params.temperatures = temperatures; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - {4}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4, 5}, // step 0 - {0}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0, 1}, // step 1 - {2}, {2, 3, 4, 5}, {2, 3}, {2, 3}, {2}, {2, 3}, // step 2 - {0}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKRepetitionPenalty) -{ - SizeType32 topK = 1; - float repetitionPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKRepetitionPenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = repetitionPenalties; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKPresencePenalty) -{ - SizeType32 topK = 1; - float presencePenalty = 1e9f; - TestSamplingParams params; - params.presencePenalties = {presencePenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKPresencePenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.presencePenalties = presencePenalties; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKFrequencyPenalty) -{ - SizeType32 topK = 1; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.frequencyPenalties = {frequencyPenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKFrequencyPenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.frequencyPenalties = frequencyPenalties; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKRepetitionPresencePenalty) -{ - SizeType32 topK = 1; - float repetitionPenalty = 1e9f; - float presencePenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.presencePenalties = {presencePenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKRepetitionPresencePenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalties}; - params.presencePenalties = {presencePenalties}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKRepetitionFrequencyPenalty) -{ - SizeType32 topK = 1; - float repetitionPenalty = 1e9f; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.frequencyPenalties = {frequencyPenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKRepetitionFrequencyPenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalties}; - params.frequencyPenalties = {frequencyPenalties}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKPresenceFrequencyPenalty) -{ - SizeType32 topK = 1; - float presencePenalty = 1e9f; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.presencePenalties = {presencePenalty}; - params.frequencyPenalties = {frequencyPenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKPresenceFrequencyPenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.presencePenalties = {presencePenalties}; - params.frequencyPenalties = {frequencyPenalties}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKFullPenalty) -{ - SizeType32 topK = 1; - float repetitionPenalty = 1e9f; - float presencePenalty = 1e9f; - float frequencyPenalty = 1e9f; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalty}; - params.presencePenalties = {presencePenalty}; - params.frequencyPenalties = {frequencyPenalty}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {1}, {1}, {1}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKFullPenaltiesBatch) -{ - SizeType32 topK = 1; - std::vector repetitionPenalties = {1e9f, 1e9f, 1.0f, 1.0f, 1.0f, 1e9f}; - std::vector presencePenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - std::vector frequencyPenalties = {1e9f, 1e9f, 0.0f, 0.0f, 0.0f, 1e9f}; - TestSamplingParams params; - params.repetitionPenalties = {repetitionPenalties}; - params.presencePenalties = {presencePenalties}; - params.frequencyPenalties = {frequencyPenalties}; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKMinLengthBatch) -{ - SizeType32 topK = 1; - std::vector minLengths = {3, 1, 1, 3, 0, 3}; - TestSamplingParams params; - params.minLengths = minLengths; - params.topKs = {topK}; - params.topPs = {1.0f}; - TokenIdType const endId = 0; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {1}, {0}, {0}, {1}, {0}, {1}, // step 1 - {2}, {0}, {0}, {2}, {0}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params, endId); -} - -TYPED_TEST(DynamicDecodeLayerTest, TopKBias) -{ - SizeType32 topK = 2; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.useBias = true; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, BadWords) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.badWords = {{{4, 0}, {2}}, {{0, 2}}, {{4, 0, 2}, {4, 0, 3, 0}}, {{3}}, {{4}, {5}}, {{0}, {3}}}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {6}, {4}, // step 0 - {1}, {0}, {0}, {0}, {0}, {1}, // step 1 - {3}, {3}, {3}, {2}, {2}, {2}, // step 2 - {0}, {0}, {1}, {0}, {0}, {1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, BadWordsNoBadWordsMode) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.badWords = {{{4, 0}, {2}}, {{0, 2}}, {{4, 0, 2}, {4, 0, 3, 0}}, {{3}}, {{4}, {5}}, {{0}, {3}}}; - params.decodingMode = tle::DecodingMode::TopK().useBanWords(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, NoRepeatNgramSize) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.badWords = {{{0}}, {{2}}, {{0}, {3}, {4, 1, 2}}, {{5}}, {{0}}, {{1}}}; - params.repeatNGramSizes = {1, 1, 2, 1, 1, 3}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {1}, {0}, {1}, {0}, {1}, {0}, // step 1 - {2}, {3}, {4}, {2}, {2}, {2}, // step 2 - {3}, {1}, {2}, {1}, {3}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, NoRepeatNgramSizeNoNgramMode) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.badWords = {{{0}}, {{2}}, {{0}, {3}, {4, 1, 2}}, {{5}}, {{0}}, {{1}}}; - params.repeatNGramSizes = {1, 1, 2, 1, 1, 3}; - params.decodingMode = tle::DecodingMode::TopK().useNoRepeatNgramSize(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {1}, {0}, {1}, {0}, {1}, {0}, // step 1 - {2}, {3}, {4}, {2}, {2}, {2}, // step 2 - {1}, {0}, {1}, {0}, {1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, NoRepeatNgramSizeNoBanTokensMode) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.badWords = {{{0}}, {{2}}, {{0}, {3}, {4, 1, 2}}, {{5}}, {{0}}, {{1}}}; - params.repeatNGramSizes = {1, 1, 2, 1, 1, 3}; - params.decodingMode = tle::DecodingMode::TopK().useBanTokens(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, NoRepeatNgramSizeMultipleRequests) -{ - this->allocateData(TestSamplingParams{}); - { - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.repeatNGramSizes = {1, 1, 2, 1, 1, 3}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {1}, {1}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } - { - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.repeatNGramSizes = {1, 1, 2, 1, 1, 3}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {1}, {1}, {0}, {1}, {1}, {0} // step 3 - }; - this->runTestImpl(expectedOutputIds, params); - } -} - -TYPED_TEST(DynamicDecodeLayerTest, StopWords) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.stopWords = {{{4, 0}, {2}}, {{0, 2}}, {{4, 0, 2}}, {{3}}, {{4}, {5}}, {{4, 0, 2, 0}}}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {2}, {2}, {2}, {0}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, StopWordsNoStopWordsMode) -{ - SizeType32 topK = 1; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {1.0f}; - params.stopWords = {{{4, 0}, {2}}, {{0, 2}}, {{4, 0, 2}}, {{3}}, {{4}, {5}}, {{4, 0, 2, 0}}}; - params.decodingMode = tle::DecodingMode::TopK().useStopWords(false); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, MedusaSimpleTest) -{ - TestSamplingParams params; - params.topKs = {1, 1, 1, 1, 1, 1}; - params.topKMedusaHeads = {{3, 1}, {1, 3}, {3, 1}, {2, 2}, {2, 2}, {1, 3}}; - params.tokensPerStep = {4, 4, 4, 4, 4, 4}; - params.maxNumMedusaHeads = 2; - // clang-format off - params.paths = {{0, 1, 2, - 0, 3, -1}, - {0, 1, -1, - 0, -1, -1}, - {0, 1, 3}, - {0, 2, 3}, - {0, 2, -1}, - {0, 3, -1}}; - // clang-format on - params.outputIds = {{4, 0, 2}, {4, 0, 2}, {4, 0, 0}, {4, 4, 2}, {4, 0, 2}, {4, 0, 2}}; - params.decodingMode = tle::DecodingMode::Medusa(); - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {2}, {4}, {4}, // step 1 - {2}, {0}, {0}, {0}, {0}, {0}, // step 2 - {2}, {2}, {0}, {2}, {2}, {2} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(DynamicDecodeLayerTest, MedusaStopWordsTest) -{ - TestSamplingParams params; - params.topKs = {1, 1, 1, 1, 1, 1}; - params.topKMedusaHeads = {{3, 1}, {1, 3}, {3, 1}, {2, 2}, {2, 2}, {1, 3}}; - params.tokensPerStep = {4, 4, 4, 4, 4, 4}; - params.maxNumMedusaHeads = 2; - // clang-format off - params.paths = {{0, 1, 2, - 0, 3, -1}, - {0, 1, -1, - 0, -1, -1}, - {0, 1, 3}, - {0, 2, 3}, - {0, 2, -1}, - {0, 3, -1}}; - // clang-format on - params.outputIds = {{4, 0, 2}, {4, 0, 2}, {4, 0, 0}, {4, 4, 2}, {4, 0, 2}, {4, 0, 2}}; - params.stopWords = {{{4, 0}}, {{0, 0}}, {{0, 2}}, {{4}, {4, 2, 0}}, {{3}}, {{4, 4, 0, 2}}}; - params.decodingMode = tle::DecodingMode::Medusa(); - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {-1}, {-1}, {-1}, // step 1 - {-1}, {-1}, {0}, {-1}, {-1}, {-1}, // step 2 - {-1}, {-1}, {-1}, {-1}, {-1}, {-1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -} // namespace tensorrt_llm::tests::layers::sampling diff --git a/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.h b/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.h deleted file mode 100644 index ca059c2266d3..000000000000 --- a/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.h +++ /dev/null @@ -1,181 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include - -#include - -#include "tensorrt_llm/layers/dynamicDecodeLayer.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -#include "tensorrt_llm/executor/types.h" - -namespace tensorrt_llm::tests::layers::sampling -{ - -struct TestSamplingParams -{ - std::vector topKs; - std::vector topPs; - std::vector temperatures; - std::vector repetitionPenalties; - std::vector presencePenalties; - std::vector frequencyPenalties; - std::vector promptIgnoreLengths; - std::vector minLengths; - std::vector decay; - std::vector minTopP; - std::vector topPResetIds; - std::vector>> badWords; - std::vector>> stopWords; - std::vector repeatNGramSizes; - bool useBias{false}; - - std::optional decodingMode; - - // Medusa setup - std::optional maxNumMedusaHeads{std::nullopt}; - std::optional>> topKMedusaHeads{std::nullopt}; - std::optional> tokensPerStep{std::nullopt}; - std::optional>> paths; - std::optional>> outputIds; -}; - -template -class DynamicDecodeLayerTest : public testing::Test -{ -private: - void SetUp() override; - - using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; - using TensorConstPtr = tensorrt_llm::runtime::ITensor::SharedConstPtr; - using BufferPtr = tensorrt_llm::runtime::IBuffer::SharedPtr; - - static uint64_t const mMaxSeed{64}; - runtime::SizeType32 const mBatchSize{6}; - runtime::SizeType32 const mMaxBatchSize{2 * mBatchSize}; - runtime::SizeType32 const mBeamWidth{1}; - runtime::SizeType32 const mBatchBeam{mBatchSize * mBeamWidth}; - runtime::SizeType32 const mVocabSize{9}; - runtime::SizeType32 const mVocabSizePadded{mVocabSize}; - - runtime::SizeType32 const mMaxInputLen{0}; // has no effect. - runtime::SizeType32 const mMaxOutputLen{4}; - runtime::SizeType32 const mMaxSeqLen{mMaxInputLen + mMaxOutputLen}; - runtime::SizeType32 const mSinkTokenLength{0}; - runtime::TokenIdType mEndId = mVocabSize; - runtime::SizeType32 mMaxTokensPerStep{1}; - runtime::SizeType32 mMaxMedusaHeads{0}; - - bool mUseLogitsVec{false}; - - TensorPtr mLogitsDevice; - TensorPtr mRuntimeLogitsHost; - TensorPtr mLogitsRefHost; - TensorPtr mContextLengthDevice; - TensorPtr mSeqLengthsDevice; - TensorPtr mFinishedDevice; - TensorPtr mFinishedSumDevice; - TensorPtr mOutputIdsDevice; - TensorPtr mNewTokens; - TensorPtr mEndIdsDevice; - TensorPtr mBatchSlots; - - TensorPtr mBadWordsLens; - TensorPtr mBadWords; - TensorPtr mBadWordsPtrs; - - TensorPtr mStopWordsLens; - TensorPtr mStopWords; - TensorPtr mStopWordsPtrs; - - TensorPtr mEmbeddingBiasHost; - TensorPtr mEmbeddingBiasDevice; - - TensorPtr mRefLogProbsHost; - TensorPtr mOutputLogProbsDevice; - TensorPtr mOutputLogProbsTiledDevice; - - TensorPtr mCumLogProbsDevice; - - // Medusa tensors - TensorPtr mPathsDevice; - TensorPtr mTreeIdsDevice; - TensorPtr mAcceptedLengths; - TensorPtr mAcceptedLengthCumSumDevice; - TensorPtr mPackedPathsDevice; - TensorPtr mMedusaLogitsDevice; - TensorPtr mNextDraftTokensDevice; - TensorPtr mTokensPerStepDevice; - - std::vector mLogitsVec; - - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - std::unique_ptr> mDecodeLayer; - std::shared_ptr mDecodingWorkspace; - - std::vector mTestLogitsInit; - - runtime::SizeType32 mMaxBadWordsLen{0}; - runtime::SizeType32 mMaxStopWordsLen{0}; - - executor::DecodingMode mDecodingMode = executor::DecodingMode::Auto(); - -private: - void allocateMedusaData(TestSamplingParams const& params); - - void setup(uint64_t seed, TestSamplingParams const& params); - - runtime::SizeType32 getMaxWordsLen(std::vector>> const& inputWords); - void initXWordsTensors(runtime::SizeType32* batchSlotsPtr, runtime::TokenIdType* wordsData, - runtime::TokenIdType** wordsPtr, runtime::SizeType32* wordsLenData, runtime::SizeType32 maxWordsLen, - std::vector>> const& inputWords); - - std::shared_ptr createInputTensors(runtime::SizeType32 step); - - std::shared_ptr createOutputTensors(); - - void batchCopy(runtime::SizeType32 step); - bool checkResult(runtime::TokenIdType* outputIds, std::vector> const& expectedIds, - runtime::SizeType32* seqLens, runtime::SizeType32 leadingDim, runtime::SizeType32 stride, - runtime::SizeType32 step, bool outputIdsTransposed = false, runtime::SizeType32 strideTransposed = 0); - - void fillRefLogits(runtime::SizeType32 const* seqLenHost, - std::vector> const& expectedOutputIds, runtime::SizeType32 step); - - void createMedusaInputs(std::shared_ptr& baseInputs); - void createMedusaOutputs(std::shared_ptr& baseOutputs); - -public: - void runTest(std::vector> const& expectedOutputIds, TestSamplingParams const& params, - runtime::TokenIdType endId = -1); - - void allocateData(TestSamplingParams const& params, runtime::TokenIdType endId = -1); - - void runTestImpl(std::vector> const& expectedOutputIds, - TestSamplingParams const& params, runtime::TokenIdType endId = -1); -}; - -typedef testing::Types FloatAndHalfTypes; - -} // namespace tensorrt_llm::tests::layers::sampling diff --git a/cpp/tests/unit_tests/layers/eagleLayerTest.cpp b/cpp/tests/unit_tests/layers/eagleLayerTest.cpp deleted file mode 100644 index 47bca93b47b4..000000000000 --- a/cpp/tests/unit_tests/layers/eagleLayerTest.cpp +++ /dev/null @@ -1,1156 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "eagleLayerTest.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/speculativeDecodingModule.h" - -#include "tensorrt_llm/common/tllmDataType.h" - -#include -#include -#include - -namespace tensorrt_llm::tests::layers -{ - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using namespace tensorrt_llm::common; - -namespace tk = tensorrt_llm::kernels; -namespace tksd = tensorrt_llm::kernels::speculative_decoding; -namespace trk = tensorrt_llm::runtime::kernels; - -///////////////////////////////////////////////////////////////////////////////////////////////////////////////// - -TokensVec EagleDummyNetwork::tokenize(std::string const& letters) const -{ - TokensVec tokens; - for (char c : letters) - { - tokens.push_back(static_cast(c)); - } - return tokens; -} - -std::string EagleDummyNetwork::detokenize(TokensVec const& tokens) const -{ - std::string letters; - for (int token : tokens) - { - letters += static_cast(token); - } - return letters; -} - -DraftTokensVec EagleDummyNetwork::draftLettersToTokens(DraftLettersVec const& draftLetters) const -{ - DraftTokensVec draftTokens(draftLetters.size()); - for (SizeType32 bi = 0; bi < draftLetters.size(); ++bi) - { - draftTokens[bi] = tokenize(draftLetters[bi]); - } - return draftTokens; -} - -SizeType32 EagleDummyNetwork::longestCommonPrefixLength(TokensVec const& a, TokensVec const& b) const -{ - SizeType32 minLength = std::min(a.size(), b.size()); - SizeType32 idx = 0; - while (idx < minLength && a[idx] == b[idx]) - { - ++idx; - } - return idx; -} - -DraftPath EagleDummyNetwork::pathFromDraftTokens( - DraftTokensVec const& tokens, SizeType32 maxDecodingTokens, SizeType32 maxPathLen) const -{ - DraftPath path(maxDecodingTokens); - for (SizeType32 pi = 0; pi < maxDecodingTokens; ++pi) - { - path[pi].resize(maxPathLen); - for (SizeType32 ti = 0; ti < maxPathLen; ++ti) - { - path[pi][ti] = -1; - } - } - SizeType32 draftPosCounter = 1; - for (SizeType32 ti = 1; ti < maxPathLen; ++ti) - { - std::unordered_map tokenPosMap; - for (SizeType32 pi = 0; pi < tokens.size(); ++pi) - { - if (tokens[pi].size() > ti - 1) - { - path[pi][0] = 0; - auto const token = tokens[pi][ti - 1]; - auto const draftPrefix = detokenize(tokens[pi]).substr(0, ti); - if (tokenPosMap.count(draftPrefix) == 0) - { - tokenPosMap[draftPrefix] = draftPosCounter++; - } - path[pi][ti] = tokenPosMap[draftPrefix]; - } - } - } - return path; -} - -TokensVec EagleDummyNetwork::flattenTokens( - DraftTokensVec const& tokens, DraftPath const& path, bool isDraftTokens) const -{ - SizeType32 maxPathIdx{-1}; - for (SizeType32 pi = 0; pi < path.size(); ++pi) - { - for (SizeType32 ti = 0; ti < path[pi].size(); ++ti) - { - auto const pathIdx = path[pi][ti]; - maxPathIdx = std::max(pathIdx, maxPathIdx); - } - } - if (!isDraftTokens) - { - maxPathIdx++; - } - TokensVec flattenedTokens(maxPathIdx); - for (SizeType32 pi = 0; pi < path.size(); ++pi) - { - for (SizeType32 ti = 0; ti < path[pi].size(); ++ti) - { - if (isDraftTokens && ti == 0) - { - continue; - } - - auto const pathIdx = path[pi][ti]; - if (pathIdx != -1) - { - if (isDraftTokens) - { - flattenedTokens[pathIdx - 1] = tokens[pi][ti - 1]; - } - else - { - flattenedTokens[pathIdx] = tokens[pi][ti]; - } - } - } - } - return flattenedTokens; -} - -std::vector>> EagleDummyNetwork::createMasks(DraftPaths const& paths) const -{ - std::vector>> masks; - for (SizeType32 bi = 0; bi < paths.size(); ++bi) - { - std::vector> localMask(paths[bi].size()); - for (SizeType32 ti = 0; ti < paths[bi].size(); ++ti) - { - localMask[ti].resize(paths[bi].size()); - } - localMask[0][0] = true; - - for (SizeType32 pi = 0; pi < paths[bi].size(); ++pi) - { - for (SizeType32 ti = 1; ti < paths[bi][pi].size(); ++ti) - { - auto const to = paths[bi][pi][ti]; - if (to == -1) - { - break; - } - localMask[to][to] = true; - for (SizeType32 fi = 0; fi < ti; ++fi) - { - auto const from = paths[bi][pi][fi]; - localMask[to][from] = true; - } - } - } - masks.push_back(localMask); - } - return masks; -} - -void EagleDummyNetwork::acceptTokens(std::vector const& predictionTokens, - DraftTokensVec const& lastDraftTokens, DraftPaths const& lastDraftPaths) -{ - TLLM_CHECK_WITH_INFO(predictionTokens.size() == lastDraftTokens.size(), - "Batch size of predictions (%d) does not match the batch size of last draft tokens (%d)", - static_cast(predictionTokens.size()), static_cast(lastDraftTokens.size())); - TLLM_CHECK_WITH_INFO(predictionTokens.size() == lastDraftPaths.size(), - "Batch size of predictions (%d) does not match the batch size of last draft paths (%d)", - static_cast(predictionTokens.size()), static_cast(lastDraftPaths.size())); - - mAcceptedTokens.resize(predictionTokens.size()); - mAcceptedLens.resize(predictionTokens.size()); - mAcceptedPathIds.resize(predictionTokens.size()); - // Needed for unit test of EagleDummyNetwork only. - if (mOutputIds.size() == 0) - { - mOutputIds.resize(lastDraftTokens.size()); - } - for (SizeType32 bi = 0; bi < lastDraftPaths.size(); ++bi) - { - SizeType32 maxMatchLen = -1; - SizeType32 maxMatchIdx = -1; - std::vector bestDraftPath; - // Find path with largest prefix shared with the predicted tokens. - for (SizeType32 pi = 0; pi < lastDraftPaths[bi].size(); ++pi) - { - TokensVec predictedPath(lastDraftPaths[bi][pi].size()); - TokensVec draftPath(lastDraftPaths[bi][pi].size()); - for (SizeType32 ti = 0; ti < lastDraftPaths[bi][pi].size(); ++ti) - { - predictedPath[ti] = predictionTokens[bi][lastDraftPaths[bi][pi][ti]]; - if (ti > 0) - { - draftPath[ti - 1] = lastDraftTokens[bi][lastDraftPaths[bi][pi][ti] - 1]; - } - } - auto const matchLen = longestCommonPrefixLength(draftPath, predictedPath); - if (matchLen > maxMatchLen) - { - maxMatchLen = matchLen; - maxMatchIdx = pi; - bestDraftPath = predictedPath; - } - } - - mAcceptedTokens[bi] = bestDraftPath; - mAcceptedLens[bi] = maxMatchLen + 1; - mAcceptedPathIds[bi] = maxMatchIdx; - // Update output ids. First draft token is already counted in outputs - mOutputIds[bi].insert(mOutputIds[bi].end(), bestDraftPath.begin(), bestDraftPath.begin() + maxMatchLen + 1); - } -} - -void EagleDummyNetwork::forward(SamplingParams const& params, std::vector const& prompts, - std::vector> const& predictionLetters, - std::vector const& nextDraftLetters, std::vector const& lastDraftLetters) -{ - mSamplingParams = params; - - TLLM_CHECK(params.getBatchSize() == nextDraftLetters.size()); - TLLM_CHECK(params.getBatchSize() == lastDraftLetters.size()); - - DraftPaths lastDraftPaths; - DraftPaths nextDraftPaths; - DraftTokensVec lastDraftTokensFlattened; - DraftTokensVec nextDraftTokensFlattened; - std::vector predictionTokensFlattened; - for (SizeType32 bi = 0; bi < params.getBatchSize(); ++bi) - { - auto const lastDraftTokens = draftLettersToTokens(lastDraftLetters[bi]); - auto const nextDraftTokens = draftLettersToTokens(nextDraftLetters[bi]); - auto const lastDraftPath - = pathFromDraftTokens(lastDraftTokens, params.getMaxDecodingTokens(), params.getMaxPathLen()); - auto const nextDraftPath - = pathFromDraftTokens(nextDraftTokens, params.getMaxDecodingTokens(), params.getMaxPathLen()); - auto const predictionTokens = draftLettersToTokens(predictionLetters[bi]); - - lastDraftPaths.push_back(lastDraftPath); - nextDraftPaths.push_back(nextDraftPath); - lastDraftTokensFlattened.push_back(flattenTokens(lastDraftTokens, lastDraftPath, /* isDraftTokens */ true)); - nextDraftTokensFlattened.push_back(flattenTokens(nextDraftTokens, nextDraftPath, /* isDraftTokens */ true)); - predictionTokensFlattened.push_back(flattenTokens(predictionTokens, lastDraftPath, /* isDraftTokens */ false)); - } - - mNextDraftTokens = nextDraftTokensFlattened; - mLastDraftTokens = lastDraftTokensFlattened; - - mNextDraftPaths = nextDraftPaths; - mLastDraftPaths = lastDraftPaths; - - mNextDraftLens.resize(params.getBatchSize()); - mLastDraftLens.resize(params.getBatchSize()); - for (SizeType32 bi = 0; bi < params.getBatchSize(); ++bi) - { - mNextDraftLens[bi] = mNextDraftTokens[bi].size(); - mLastDraftLens[bi] = mLastDraftTokens[bi].size(); - } - - std::vector predictionTokens; - for (SizeType32 bi = 0; bi < predictionLetters.size(); ++bi) - { - mPrompts.push_back(tokenize(prompts[bi])); - } - - mOutputIds = mPrompts; - - acceptTokens(predictionTokensFlattened, mLastDraftTokens, lastDraftPaths); - - mMasks = createMasks(mNextDraftPaths); -} - -TEST(EagleDummyNetworkTest, tokenizeTest) -{ - EagleDummyNetwork network; - - { - auto tokens = network.tokenize("hello world"); - EXPECT_EQ(tokens, std::vector({104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100})); - } - { - DraftLettersVec lettersVec = {{"hello world"}, {"world"}}; - auto draftTokens = network.draftLettersToTokens(lettersVec); - ASSERT_EQ(draftTokens.size(), 2); - ASSERT_EQ(draftTokens[0].size(), 11); - ASSERT_EQ(draftTokens[1].size(), 5); - EXPECT_EQ(draftTokens[0], std::vector({104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100})); - EXPECT_EQ(draftTokens[1], std::vector({119, 111, 114, 108, 100})); - } -} - -TEST(EagleDummyNetworkTest, detokenizeTest) -{ - EagleDummyNetwork network; - - { - auto letters - = network.detokenize(std::vector({104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100})); - EXPECT_EQ(letters, "hello world"); - } -} - -TEST(EagleDummyNetworkTest, longestCommonPrefixLengthTest) -{ - EagleDummyNetwork network; - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {1, 2}), 2); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {1, 2, 3}), 3); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {1, 5, 6}), 1); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {2, 5, 6}), 0); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {}), 0); -} - -TEST(EagleDummyNetworkTest, pathFromDraftTokensTest) -{ - EagleDummyNetwork network; - { - SizeType32 const maxDecodingTokens = 5; - SizeType32 const maxPathLen = 4; - DraftTokensVec draftTokens = {{1, 4, 8}, {1, 5}, {2, 6, 9}, {2, 7}, {3}}; - auto const paths = network.pathFromDraftTokens(draftTokens, maxDecodingTokens, maxPathLen); - ASSERT_EQ(paths.size(), maxDecodingTokens); - for (SizeType32 pi = 0; pi < maxDecodingTokens; ++pi) - { - ASSERT_EQ(paths[pi].size(), maxPathLen); - if (pi < draftTokens.size()) - { - for (SizeType32 ti = 0; ti < maxPathLen; ++ti) - { - if (ti == 0) - { - EXPECT_EQ(paths[pi][ti], 0); - } - else if (ti - 1 < draftTokens[pi].size()) - { - EXPECT_EQ(paths[pi][ti], draftTokens[pi][ti - 1]); - } - else - { - EXPECT_EQ(paths[pi][ti], -1); - } - } - } - else - { - for (SizeType32 ti = 0; ti < maxPathLen; ++ti) - { - EXPECT_EQ(paths[pi][ti], -1); - } - } - } - } -} - -TEST(EagleDummyNetworkTest, flattenedTokensTest) -{ - { - EagleDummyNetwork network; - DraftTokensVec draftTokens = {{1, 4, 8}, {1, 5}, {2, 6, 9}, {2, 7}, {3}}; - DraftPath path = {{0, 1, 4, 8}, {0, 1, 5, -1}, {0, 2, 6, 9}, {0, 2, 7, -1}, {0, 3, -1, -1}, {-1, -1, -1, -1}, - {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}; - - auto const flattenTokens = network.flattenTokens(draftTokens, path, /* isDraftTokens*/ true); - EXPECT_EQ(flattenTokens, TokensVec({1, 2, 3, 4, 5, 6, 7, 8, 9})); - } - { - EagleDummyNetwork network; - DraftTokensVec predictionTokens = {{0, 1, 4, 8}, {0, 1, 5}, {0, 2, 6, 9}, {0, 2, 7}, {0, 3}}; - DraftPath path = {{0, 1, 4, 8}, {0, 1, 5, -1}, {0, 2, 6, 9}, {0, 2, 7, -1}, {0, 3, -1, -1}, {-1, -1, -1, -1}, - {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}; - - auto const flattenTokens = network.flattenTokens(predictionTokens, path, /* isDraftTokens*/ false); - EXPECT_EQ(flattenTokens, TokensVec({0, 1, 2, 3, 4, 5, 6, 7, 8, 9})); - } -} - -TEST(EagleDummyNetworkTest, createMasksTest) -{ - { - EagleDummyNetwork network; - DraftPaths paths = {{{0, 1, -1, -1}, {-1, -1, -1, -1}}}; - - auto const mask = network.createMasks(paths); - std::vector>> refMask = {{{true, false}, {true, true}}}; - EXPECT_EQ(mask, refMask); - } - { - EagleDummyNetwork network; - DraftPaths paths = {{{0, 1, 4, 8}, {0, 1, 5, -1}, {0, 2, 6, 9}, {0, 2, 7, -1}, {0, 3, -1, -1}, {-1, -1, -1, -1}, - {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}}; - - auto const mask = network.createMasks(paths); - std::vector>> refMask - = {{{true, false, false, false, false, false, false, false, false, false}, - {true, true, false, false, false, false, false, false, false, false}, - {true, false, true, false, false, false, false, false, false, false}, - {true, false, false, true, false, false, false, false, false, false}, - {true, true, false, false, true, false, false, false, false, false}, - {true, true, false, false, false, true, false, false, false, false}, - {true, false, true, false, false, false, true, false, false, false}, - {true, false, true, false, false, false, false, true, false, false}, - {true, true, false, false, true, false, false, false, true, false}, - {true, false, true, false, false, false, true, false, false, true}}}; - EXPECT_EQ(mask, refMask); - } - { - EagleDummyNetwork network; - DraftPaths paths = {{{0, 1, 3}, {0, 2, -1}, {-1, -1, -1}, {-1, -1, -1}, {-1, -1, -1}}, - {{0, 1, 3}, {0, 2, 4}, {-1, -1, -1}, {-1, -1, -1}, {-1, -1, -1}}}; - - auto const mask = network.createMasks(paths); - std::vector>> refMask = { - {{true, false, false, false, false}, {true, true, false, false, false}, {true, false, true, false, false}, - {true, true, false, true, false}, {false, false, false, false, false}}, - {{true, false, false, false, false}, {true, true, false, false, false}, {true, false, true, false, false}, - {true, true, false, true, false}, {true, false, true, false, true}}}; - EXPECT_EQ(mask, refMask); - } -} - -TEST(EagleDummyNetworkTest, acceptTokensTest) -{ - { - EagleDummyNetwork network; - SizeType32 const batchSize{1}; - SizeType32 const maxDecodingTokens{10}; - SizeType32 const maxPathLen{4}; - std::vector predictionLetters = {{"howe", "hoc", "hecl", "hea", "hu"}}; - std::vector lastDraftLetters = {{"how", "he", "wow", "we", "a"}}; - DraftPaths lastDraftPaths; - DraftTokensVec lastDraftTokensFlattened; - std::vector predictionTokensFlattened; - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto const lastDraftTokens = network.draftLettersToTokens(lastDraftLetters[bi]); - auto const lastDraftPath = network.pathFromDraftTokens(lastDraftTokens, maxDecodingTokens, maxPathLen); - auto const predictionTokens = network.draftLettersToTokens(predictionLetters[bi]); - lastDraftPaths.push_back(lastDraftPath); - lastDraftTokensFlattened.push_back( - network.flattenTokens(lastDraftTokens, lastDraftPath, /* isDraftTokens */ true)); - predictionTokensFlattened.push_back( - network.flattenTokens(predictionTokens, lastDraftPath, /* isDraftTokens */ false)); - } - - network.acceptTokens(predictionTokensFlattened, lastDraftTokensFlattened, lastDraftPaths); - - auto acceptedLens = network.getAcceptedLens(); - auto acceptedPathIds = network.getAcceptedPathIds(); - auto outputIds = network.getOutputIds(); - - ASSERT_EQ(acceptedLens.size(), 1); - ASSERT_EQ(acceptedPathIds.size(), 1); - ASSERT_EQ(outputIds.size(), 1); - EXPECT_EQ(acceptedLens[0], 4); - EXPECT_EQ(acceptedPathIds[0], 0); - EXPECT_EQ(network.detokenize(outputIds[0]), "howe"); - } - - { - EagleDummyNetwork network; - SizeType32 const batchSize{2}; - SizeType32 const maxDecodingTokens{10}; - SizeType32 const maxPathLen{4}; - std::vector predictionLetters - = {{"howe", "hoc", "hecl", "hea", "hu"}, {"bcde", "bcdc", "bca", "bcc", "bo"}}; - std::vector lastDraftLetters - = {{"how", "he", "wow", "we", "a"}, {"inc", "inf", "ir", "im", "b"}}; - DraftPaths lastDraftPaths; - DraftTokensVec lastDraftTokensFlattened; - std::vector predictionTokensFlattened; - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - auto const lastDraftTokens = network.draftLettersToTokens(lastDraftLetters[bi]); - auto const lastDraftPath = network.pathFromDraftTokens(lastDraftTokens, maxDecodingTokens, maxPathLen); - auto const predictionTokens = network.draftLettersToTokens(predictionLetters[bi]); - lastDraftPaths.push_back(lastDraftPath); - lastDraftTokensFlattened.push_back( - network.flattenTokens(lastDraftTokens, lastDraftPath, /* isDraftTokens */ true)); - predictionTokensFlattened.push_back( - network.flattenTokens(predictionTokens, lastDraftPath, /* isDraftTokens */ false)); - } - - network.acceptTokens(predictionTokensFlattened, lastDraftTokensFlattened, lastDraftPaths); - - auto acceptedLens = network.getAcceptedLens(); - auto acceptedPathIds = network.getAcceptedPathIds(); - auto outputIds = network.getOutputIds(); - - ASSERT_EQ(acceptedLens.size(), 2); - ASSERT_EQ(acceptedPathIds.size(), 2); - ASSERT_EQ(outputIds.size(), 2); - EXPECT_EQ(acceptedLens[0], 4); - EXPECT_EQ(acceptedLens[1], 2); - EXPECT_EQ(acceptedPathIds[0], 0); - EXPECT_EQ(acceptedPathIds[1], 4); - EXPECT_EQ(network.detokenize(outputIds[0]), "howe"); - EXPECT_EQ(network.detokenize(outputIds[1]), "bo"); - } -} - -///////////////////////////////////////////////////////////////////////////////////////////////////////////////// - -template -void EagleDecodingLayerTest::SetUp() -{ - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); -} - -template -void EagleDecodingLayerTest::allocateBuffers() -{ - auto speculativeDecodingModule = std::make_shared(mSamplingParams.getMaxDraftPathLen(), - mSamplingParams.getMaxDecodingDraftTokens(), mSamplingParams.getMaxDecodingTokens()); - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain(mSamplingParams.getMaxBatchSize(), 1, - mSamplingParams.getVocabSize(), mSamplingParams.getVocabSize(), speculativeDecodingModule); - - mEagleLayer = std::make_shared>(decodingDomain, mBufferManager); - - // outputs - mOutputIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxSeqLen()}), - tensorrt_llm::DataType::kINT32); - - mSeqLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mOutputNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mOutputUnpackedNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mAcceptedLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextPosIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - - mPrevDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextGenerationLengths - = mBufferManager->gpu(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextGenerationLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mAcceptedLengthCumSum = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() + 1}), tensorrt_llm::DataType::kINT32); - - mPathsOffsets = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() * mSamplingParams.getMaxDraftPathLen()}), - tensorrt_llm::DataType::kINT32); - - mPackedMasks = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), - static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))}), - tensorrt_llm::DataType::kINT32); - - mRandomDataSample = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); - - mRandomDataValidation = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kFLOAT); - - mOutputTemperatures = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); - - mOutputNextDraftPaths - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize(), - mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mEagleNetCtxRequestTypesHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mEagleNetCtxContextLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mEagleNetCtxPastKeyValueLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mEagleNetGenRequestTypesHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mEagleNetGenContextLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mEagleNetGenPastKeyValueLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - // inputs - mBatchSlots = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mEndIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mInputNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mInputNextDraftLens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mInputNextDraftPaths = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mInputLastDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mInputLastDraftLens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mInputLastDraftPaths = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mInputAcceptedTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mInputAcceptedLens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mInputAcceptedPathIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mChunkedContextNextTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mDecodingWorkspace = std::make_shared(mBufferManager, decodingDomain, - TRTDataType::value, mSamplingParams.getMaxBatchSize() * sizeof(curandState_t)); -} - -template -void EagleDecodingLayerTest::setup() -{ - // outputs - trk::invokeFill(*mOutputIds, TokenIdType{-1}, *mStream); - trk::invokeFill(*mSeqLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mOutputNextDraftTokens, TokenIdType{-1}, *mStream); - trk::invokeFill(*mOutputUnpackedNextDraftTokens, TokenIdType{-1}, *mStream); - trk::invokeFill(*mAcceptedLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mNextPosIds, SizeType32{0}, *mStream); - trk::invokeFill(*mPrevDraftLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mNextDraftLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mNextGenerationLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mNextGenerationLengthsHost, SizeType32{0}, *mStream); - trk::invokeFill(*mAcceptedLengthCumSum, SizeType32{-1}, *mStream); - trk::invokeFill(*mPathsOffsets, SizeType32{0}, *mStream); - trk::invokeFill(*mPackedMasks, SizeType32{0}, *mStream); - trk::invokeFill(*mEndIds, TokenIdType{-1}, *mStream); - trk::invokeFill(*mRandomDataSample, float{0}, *mStream); - trk::invokeFill(*mRandomDataValidation, float{0}, *mStream); - trk::invokeFill(*mOutputTemperatures, float{0}, *mStream); - trk::invokeFill(*mOutputNextDraftPaths, SizeType32{0}, *mStream); - trk::invokeFill(*mChunkedContextNextTokens, SizeType32{-1}, *mStream); - - std::mt19937 gen(42); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - batchSlotsPtr[bi] = 2 * bi; - } - - auto setupParams = std::make_shared(); - mRandomSeeds = std::vector(mSamplingParams.getBatchSize()); - mTemperatures = std::vector(mSamplingParams.getBatchSize()); - - std::mt19937 generator(42); - std::uniform_int_distribution seedDistr(1, 1000); - std::uniform_real_distribution temperatureDistr(0.001f, 1.f); - std::generate( - mRandomSeeds.begin(), mRandomSeeds.end(), [&generator, &seedDistr]() { return seedDistr(generator); }); - std::generate(mTemperatures.begin(), mTemperatures.end(), - [&generator, &temperatureDistr]() { return temperatureDistr(generator); }); - setupParams->randomSeed = mRandomSeeds; - setupParams->temperature = mTemperatures; - setupParams->randomDataSample = mRandomDataSample; - setupParams->temperatures = mOutputTemperatures; - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mEagleLayer->setup(mSamplingParams.getBatchSize(), 1, mBatchSlots, setupParams, mDecodingWorkspace); - - mStream->synchronize(); - - mInputAcceptedLens = mBufferManager->copyFrom(mNetwork.getAcceptedLens(), - ITensor::makeShape({mSamplingParams.getBatchSize()}), runtime::MemoryType::kPINNEDPOOL); - mInputAcceptedPathIds = mBufferManager->copyFrom(mNetwork.getAcceptedPathIds(), - ITensor::makeShape({mSamplingParams.getBatchSize()}), runtime::MemoryType::kPINNEDPOOL); - - auto const nextDraftTokens = mNetwork.getNextDraftTokens(); - auto const lastDraftTokens = mNetwork.getLastDraftTokens(); - auto const nextDraftPaths = mNetwork.getNextDraftPaths(); - auto const lastDraftPaths = mNetwork.getLastDraftPaths(); - auto const nextDraftLens = mNetwork.getNextDraftLens(); - auto const lastDraftLens = mNetwork.getLastDraftLens(); - auto const acceptedTokens = mNetwork.getAcceptedTokens(); - auto sequenceLength = BufferRange(*mSeqLengths); - auto inputNextDraftTokensRange = BufferRange(*mInputNextDraftTokens); - auto inputLastDraftTokensRange = BufferRange(*mInputLastDraftTokens); - auto inputNextDraftPathsRange = BufferRange(*mInputNextDraftPaths); - auto inputLastDraftPathsRange = BufferRange(*mInputLastDraftPaths); - auto inputNextDraftLensRange = BufferRange(*mInputNextDraftLens); - auto inputLastDraftLensRange = BufferRange(*mInputLastDraftLens); - auto inputAcceptedTokensRange = BufferRange(*mInputAcceptedTokens); - - auto outputIds = BufferRange(*mOutputIds); - auto prompts = mNetwork.getPrompts(); - for (SizeType32 bi = 0; bi < nextDraftTokens.size(); ++bi) - { - for (SizeType32 ti = 0; ti < nextDraftTokens[bi].size(); ++ti) - { - auto idx = flat_index2(bi, ti, mSamplingParams.getMaxDecodingDraftTokens()); - inputNextDraftTokensRange[idx] = nextDraftTokens[bi][ti]; - } - for (SizeType32 ti = 0; ti < lastDraftTokens[bi].size(); ++ti) - { - auto idx = flat_index2(bi, ti, mSamplingParams.getMaxDecodingDraftTokens()); - inputLastDraftTokensRange[idx] = lastDraftTokens[bi][ti]; - } - for (SizeType32 pi = 0; pi < nextDraftPaths[bi].size(); ++pi) - { - for (SizeType32 ti = 0; ti < nextDraftPaths[bi][pi].size(); ++ti) - { - auto idx - = flat_index3(bi, pi, ti, mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()); - inputNextDraftPathsRange[idx] = nextDraftPaths[bi][pi][ti]; - } - } - for (SizeType32 pi = 0; pi < lastDraftPaths[bi].size(); ++pi) - { - for (SizeType32 ti = 0; ti < lastDraftPaths[bi][pi].size(); ++ti) - { - auto idx - = flat_index3(bi, pi, ti, mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()); - inputLastDraftPathsRange[idx] = lastDraftPaths[bi][pi][ti]; - } - } - inputNextDraftLensRange[bi] = nextDraftLens[bi]; - inputLastDraftLensRange[bi] = lastDraftLens[bi]; - sequenceLength[batchSlotsPtr[bi]] = prompts[bi].size(); - } - for (SizeType32 bi = 0; bi < acceptedTokens.size(); ++bi) - { - for (SizeType32 ti = 0; ti < acceptedTokens[bi].size(); ++ti) - { - auto idx = flat_index2(bi, ti, mSamplingParams.getMaxPathLen()); - inputAcceptedTokensRange[idx] = acceptedTokens[bi][ti]; - } - } -} - -template -std::shared_ptr EagleDecodingLayerTest::createInputTensors() -{ - auto forwardParams - = std::make_shared(mEndIds, mBatchSlots, mSamplingParams.getBatchSize(), mInputNextDraftTokens, - mInputNextDraftLens, mInputNextDraftPaths, mInputLastDraftTokens, mInputLastDraftLens, mInputLastDraftPaths, - mInputAcceptedTokens, mInputAcceptedLens, mInputAcceptedPathIds, mChunkedContextNextTokens, mBatchSlots); - - return forwardParams; -} - -template -std::shared_ptr EagleDecodingLayerTest::createOutputTensors() -{ - auto outputParams = std::make_shared(mOutputIds); - - outputParams->sequenceLength = mSeqLengths; - - outputParams->unpackedNextDraftTokens = mOutputUnpackedNextDraftTokens; - - outputParams->nextDraftTokens = mOutputNextDraftTokens; - - outputParams->numNewTokens = mAcceptedLengths; - - outputParams->nextDraftPosIds = mNextPosIds; - - outputParams->prevDraftLengths = mPrevDraftLengths; - - outputParams->nextDraftLengths = mNextDraftLengths; - - outputParams->generationLengths = mNextGenerationLengths; - - outputParams->generationLengthsHost = mNextGenerationLengthsHost; - - outputParams->numNewTokensCumSum = mAcceptedLengthCumSum; - - outputParams->pathsOffsets = mPathsOffsets; - - outputParams->packedMasks = mPackedMasks; - - outputParams->randomDataSample = mRandomDataSample; - - outputParams->randomDataValidation = mRandomDataValidation; - - outputParams->temperatures = mOutputTemperatures; - - outputParams->nextDraftPaths = mOutputNextDraftPaths; - - outputParams->eagleNetCtxRequestTypesHost = mEagleNetCtxRequestTypesHost; - - outputParams->eagleNetCtxContextLengthsHost = mEagleNetCtxContextLengthsHost; - - outputParams->eagleNetCtxPastKeyValueLengthsHost = mEagleNetCtxPastKeyValueLengthsHost; - - outputParams->eagleNetGenRequestTypesHost = mEagleNetGenRequestTypesHost; - - outputParams->eagleNetGenContextLengthsHost = mEagleNetGenContextLengthsHost; - - outputParams->eagleNetGenPastKeyValueLengthsHost = mEagleNetGenPastKeyValueLengthsHost; - - return outputParams; -} - -std::vector boolArrayToBitmask(std::vector::iterator boolIterator, size_t pathLen) -{ - std::vector bitmask(divUp(pathLen, 32)); - for (size_t bi = 0; bi < pathLen; ++bi) - { - auto slice = bi / 32; - if (boolIterator[bi]) - { - bitmask[slice] |= (1 << (bi % 32)); - } - } - return bitmask; -} - -template -void EagleDecodingLayerTest::checkLayerResult() -{ - auto const batchSlots = BufferRange(*mBatchSlots); - - // Check generated random data - { - auto const randomDataSample = BufferRange(*mRandomDataSample); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - // Check that all fields are filled with non zero data - EXPECT_NE(randomDataSample[batchSlot], float{0}) << " bi: " << bi; - } - } - - // Check masks - { - auto const randomDataValidation = BufferRange(*mRandomDataValidation); - auto const packedMasks = BufferRange(*mPackedMasks); - auto masks = mNetwork.getNextMasks(); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDecodingTokens(); ++ti) - { - auto const batchSlot = batchSlots[bi]; - auto const bitmask = boolArrayToBitmask(masks[bi][ti].begin(), mSamplingParams.getMaxDecodingTokens()); - - EXPECT_NE(randomDataValidation[batchSlot * mSamplingParams.getMaxDecodingTokens() + ti], float{0}) - << " bi: " << bi; - - for (SizeType32 mi = 0; mi < bitmask.size(); ++mi) - { - auto const packedMaskIdx = flat_index3(batchSlot, ti, mi, mSamplingParams.getMaxDecodingTokens(), - static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))); - EXPECT_EQ(bitmask[mi], packedMasks[packedMaskIdx]) << " bi: " << bi << " ti: " << ti; - } - } - } - } - - // Check accepted tokens - auto const outputIds = BufferRange(*mOutputIds); - auto const refOutputIds = mNetwork.getOutputIds(); - auto const promptIds = mNetwork.getPrompts(); - auto const seqLenghts = BufferRange(*mSeqLengths); - auto const acceptedLengths = BufferRange(*mAcceptedLengths); - auto const inputAcceptedLens = BufferRange(*mInputAcceptedLens); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - // Check accepted length. - EXPECT_EQ(inputAcceptedLens[bi], acceptedLengths[batchSlot]) << " bi:" << bi; - // Updated seq length is prompt length and newly accepted tokens. - EXPECT_EQ(seqLenghts[batchSlot], promptIds[bi].size() + acceptedLengths[batchSlot]) << " bi: " << bi; - // Check that output ids contains accepted tokens. - for (SizeType32 ti = promptIds[bi].size(); ti < acceptedLengths[batchSlot]; ++ti) - { - EXPECT_EQ(outputIds[batchSlot * mSamplingParams.getMaxSeqLen() + ti], refOutputIds[bi][ti]) - << " bi: " << bi << " ti: " << ti; - } - } - - // Check new draft tokens - { - auto const outputNextDraftTokens = BufferRange(*mOutputNextDraftTokens); - auto const outputUnpackedNextDraftTokens = BufferRange(*mOutputUnpackedNextDraftTokens); - auto const nextDraftLens = mNetwork.getNextDraftLens(); - auto const prevDraftLens = mNetwork.getLastDraftLens(); - auto const nextDraftTokens = mNetwork.getNextDraftTokens(); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - auto const nextDraftLen = nextDraftLens[bi]; - auto const prevDraftLen = prevDraftLens[bi]; - // Check draft tokens for the next iteration. - for (SizeType32 ti = 0; ti < nextDraftLen; ++ti) - { - auto const idx = flat_index2(batchSlot, ti, mSamplingParams.getMaxDecodingDraftTokens()); - EXPECT_EQ(outputNextDraftTokens[idx], nextDraftTokens[bi][ti]) << " bi: " << bi << " ti: " << ti; - EXPECT_EQ(outputUnpackedNextDraftTokens[idx], nextDraftTokens[bi][ti]) - << " bi: " << bi << " ti: " << ti; - } - // Check length of the draft tokens. - EXPECT_EQ(BufferRange(*mNextGenerationLengthsHost)[batchSlot], nextDraftLen + 1) - << " bi: " << bi; - EXPECT_EQ(BufferRange(*mNextDraftLengths)[batchSlot], nextDraftLen) << " bi: " << bi; - EXPECT_EQ(BufferRange(*mPrevDraftLengths)[batchSlot], prevDraftLen) << " bi: " << bi; - - for (SizeType32 pi = 0; pi < mSamplingParams.getMaxDecodingTokens(); ++pi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxPathLen(); ++ti) - { - auto const idx = flat_index3( - bi, pi, ti, mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()); - auto const idxSlot = flat_index3( - batchSlot, pi, ti, mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()); - EXPECT_EQ(BufferRange(*mOutputNextDraftPaths)[idxSlot], - BufferRange(*mInputNextDraftPaths)[idx]) - << " bi: " << bi << " pi:" << pi << " ti: " << ti; - } - } - } - } - - // Check position ids - { - auto const nextPosIds = BufferRange(*mNextPosIds); - auto const nextDraftPaths = mNetwork.getNextDraftPaths(); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - // Check pos ids for the next iteration. - for (SizeType32 pi = 0; pi < mSamplingParams.getMaxDecodingTokens(); ++pi) - { - for (SizeType32 li = 0; li < mSamplingParams.getMaxPathLen(); ++li) - { - auto const pathIdx = nextDraftPaths[bi][pi][li]; - auto const idx = flat_index2(batchSlot, pathIdx, mSamplingParams.getMaxDecodingTokens()); - if (pathIdx != -1) - { - EXPECT_EQ(nextPosIds[idx], li) << " bi: " << bi << " pi: " << pi << " li: " << li; - } - } - } - } - } - - // Check accumulated cum sum and paths offsets - { - auto const accumulatedCumSum = BufferRange(*mAcceptedLengthCumSum); - auto const pathsOffsets = BufferRange(*mPathsOffsets); - auto const acceptedLengths = BufferRange(*mAcceptedLengths); - auto const inputAcceptedPathIds = BufferRange(*mInputAcceptedPathIds); - auto const lastDraftPaths = mNetwork.getLastDraftPaths(); - SizeType32 sum = 0; - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ(sum, accumulatedCumSum[bi]) << "bi: " << bi; - auto const acceptedLength = acceptedLengths[batchSlot] - 1; - for (SizeType32 ti = 0; ti < acceptedLength; ++ti) - { - EXPECT_EQ(pathsOffsets[sum + ti], lastDraftPaths[bi][inputAcceptedPathIds[bi]][ti + 1] - 1) - << "bi: " << bi << " ti: " << ti; - } - sum += acceptedLength; - } - EXPECT_EQ(sum, accumulatedCumSum[mSamplingParams.getBatchSize()]); - } - - // Check temperature - { - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ(BufferRange(*mOutputTemperatures)[batchSlot], static_cast(mTemperatures[bi])) - << " bi: " << bi; - } - } - - // Check EagleNet host buffers - { - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ(BufferRange(*mEagleNetCtxRequestTypesHost)[batchSlot], 0) << " bi: " << bi; - EXPECT_EQ(BufferRange(*mEagleNetGenRequestTypesHost)[batchSlot], 1) << " bi: " << bi; - - EXPECT_EQ( - BufferRange(*mEagleNetCtxContextLengthsHost)[batchSlot], mSamplingParams.getMaxPathLen()) - << " bi: " << bi; - EXPECT_EQ(BufferRange(*mEagleNetGenContextLengthsHost)[batchSlot], - seqLenghts[batchSlot] + mSamplingParams.getMaxPathLen()) - << " bi: " << bi; - - EXPECT_EQ(BufferRange(*mEagleNetCtxPastKeyValueLengthsHost)[batchSlot], - seqLenghts[batchSlot] + mSamplingParams.getMaxPathLen()) - << " bi: " << bi; - EXPECT_EQ(BufferRange(*mEagleNetGenPastKeyValueLengthsHost)[batchSlot], - seqLenghts[batchSlot] + mSamplingParams.getMaxPathLen() - 1) - << " bi: " << bi; - } - } -} - -template -void EagleDecodingLayerTest::runTest(std::vector const& prompts, - std::vector const& predictions, std::vector const& nextDraftLetters, - std::vector const& lastDraftLetters, SamplingParams& params) -{ - mSamplingParams = params; - - mNetwork.forward(params, prompts, predictions, nextDraftLetters, lastDraftLetters); - - allocateBuffers(); - - setup(); - - auto inputTensors = createInputTensors(); - auto outputTensors = createOutputTensors(); - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mEagleLayer->forwardAsync(outputTensors, inputTensors, mDecodingWorkspace); - - mStream->synchronize(); - - checkLayerResult(); -} - -TYPED_TEST_SUITE(EagleDecodingLayerTest, FloatAndHalfTypes); - -TYPED_TEST(EagleDecodingLayerTest, IOSamePathsBs1) -{ - SamplingParams params; - - params.setBatchSize(1); - params.setMaxPathLen(4); - params.setMaxDecodingTokens(10); - - std::vector prompts = {"Hi mate, "}; - std::vector predictionLetters = {{"how ", "hoc", "hecl", "hea", "hu"}}; - std::vector lastDraftLetters = {{"how", "he", "wow", "we", "a"}}; - std::vector nextDraftLetters = {{"are", "ap", "cre", "co", "i"}}; - - this->runTest(prompts, predictionLetters, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(EagleDecodingLayerTest, IODifferentPathsBs1) -{ - SamplingParams params; - - params.setBatchSize(1); - params.setMaxPathLen(4); - params.setMaxDecodingTokens(10); - - std::vector prompts = {"Hi mate, "}; - std::vector predictionLetters = {{"how ", "hoc", "hecl", "hea", "hu"}}; - std::vector lastDraftLetters = {{"how", "he", "wow", "we", "a"}}; - std::vector nextDraftLetters = {{"are", "is", "imp", "do"}}; - - this->runTest(prompts, predictionLetters, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(EagleDecodingLayerTest, IODifferentPathsNoDraftAcceptedBs1) -{ - SamplingParams params; - - params.setBatchSize(1); - params.setMaxPathLen(4); - params.setMaxDecodingTokens(10); - - std::vector prompts = {"Hi mate, "}; - std::vector predictionLetters = {{"how ", "hoc", "hecl", "hea", "hu"}}; - std::vector lastDraftLetters = {{"my", "I'd", "wow", "we", "a"}}; - std::vector nextDraftLetters = {{"are", "ap", "cre", "co", "i"}}; - - this->runTest(prompts, predictionLetters, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(EagleDecodingLayerTest, IODifferentPathsBs2) -{ - SamplingParams params; - - params.setBatchSize(2); - params.setMaxPathLen(4); - params.setMaxDecodingTokens(10); - - std::vector prompts = {"Hi mate, ", "Let's go "}; - std::vector predictionLetters - = {{"how ", "hoc", "hecl", "hea", "hu"}, {"bcde", "bcdc", "bca", "bcc", "bo"}}; - std::vector lastDraftLetters = {{"how", "he", "wow", "we", "a"}, {"inc", "inf", "ir", "im", "b"}}; - std::vector nextDraftLetters = {{"are", "is", "imp", "do"}, {"wli", "mbi", "ard"}}; - - this->runTest(prompts, predictionLetters, nextDraftLetters, lastDraftLetters, params); -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/eagleLayerTest.h b/cpp/tests/unit_tests/layers/eagleLayerTest.h deleted file mode 100644 index c9b7350ba36d..000000000000 --- a/cpp/tests/unit_tests/layers/eagleLayerTest.h +++ /dev/null @@ -1,303 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "tensorrt_llm/layers/eagleDecodingLayer.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -#include - -#include - -namespace tensorrt_llm::tests::layers -{ - -class SamplingParams -{ -public: - SamplingParams() {} - - inline void setBatchSize(runtime::SizeType32 batchSize) - { - mBatchSize = batchSize; - } - - inline void setMaxPathLen(runtime::SizeType32 maxPathLen) - { - mMaxPathLen = maxPathLen; - } - - inline void setMaxDecodingTokens(runtime::SizeType32 maxDecodingTokens) - { - mMaxDecodingTokens = maxDecodingTokens; - } - - [[nodiscard]] inline runtime::SizeType32 getBatchSize() const - { - return mBatchSize; - } - - [[nodiscard]] inline runtime::SizeType32 getVocabSize() const - { - return mVocabSize; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxBatchSize() const - { - return 2 * getBatchSize(); - } - - [[nodiscard]] inline runtime::SizeType32 getMaxPathLen() const - { - return mMaxPathLen; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxDraftPathLen() const - { - return getMaxPathLen() - 1; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxDecodingTokens() const - { - return mMaxDecodingTokens; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxDecodingDraftTokens() const - { - return getMaxDecodingTokens() - 1; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxSeqLen() const - { - return getMaxDecodingTokens() * 2; - } - - [[nodiscard]] inline runtime::TokenIdType getPadId() const - { - return mPadId; - } - -private: - runtime::SizeType32 mBatchSize{6}; - runtime::SizeType32 mMaxPathLen{4}; - runtime::SizeType32 mMaxDecodingTokens{32}; - runtime::SizeType32 mVocabSize{256}; - runtime::TokenIdType mPadId{-1}; -}; - -using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; -using BufferPtr = tensorrt_llm::runtime::IBuffer::SharedPtr; -using SizeType32 = tensorrt_llm::runtime::SizeType32; -using TokenIdType = tensorrt_llm::runtime::TokenIdType; - -using TokensVec = std::vector; -using DraftLettersVec = std::vector; -using DraftTokensVec = std::vector; -using DraftPath = std::vector>; -using DraftPaths = std::vector; - -class EagleDummyNetwork -{ -public: - void forward(SamplingParams const& params, std::vector const& prompts, - std::vector> const& predictionLetters, - std::vector const& nextDraftLetters, std::vector const& lastDraftLetters); - - TokensVec tokenize(std::string const& letters) const; - - std::string detokenize(TokensVec const& tokens) const; - - SizeType32 longestCommonPrefixLength(TokensVec const& a, TokensVec const& b) const; - - DraftTokensVec draftLettersToTokens(DraftLettersVec const& draftLetters) const; - - DraftPath pathFromDraftTokens( - DraftTokensVec const& tokens, SizeType32 maxDecodingTokens, SizeType32 maxPathLen) const; - - TokensVec flattenTokens(DraftTokensVec const& tokens, DraftPath const& path, bool isDraftTokens) const; - - void acceptTokens(std::vector const& predictionTokens, DraftTokensVec const& lastDraftTokens, - DraftPaths const& lastDraftPaths); - - std::vector>> createMasks(DraftPaths const& paths) const; - - void setSamplingParams(SamplingParams const& params) - { - mSamplingParams = params; - } - - std::vector getPrompts() const - { - return mPrompts; - } - - std::vector getOutputIds() const - { - return mOutputIds; - } - - DraftTokensVec getNextDraftTokens() const - { - return mNextDraftTokens; - } - - std::vector getNextDraftLens() const - { - return mNextDraftLens; - } - - DraftPaths getNextDraftPaths() const - { - return mNextDraftPaths; - } - - DraftTokensVec getLastDraftTokens() const - { - return mLastDraftTokens; - } - - std::vector getLastDraftLens() const - { - return mLastDraftLens; - } - - DraftPaths getLastDraftPaths() const - { - return mLastDraftPaths; - } - - std::vector getAcceptedTokens() const - { - return mAcceptedTokens; - } - - std::vector getAcceptedLens() const - { - return mAcceptedLens; - } - - std::vector getAcceptedPathIds() const - { - return mAcceptedPathIds; - } - - std::vector>> getNextMasks() const - { - return mMasks; - } - -private: - SamplingParams mSamplingParams; - - std::vector mPrompts; - std::vector mOutputIds; - - DraftTokensVec mNextDraftTokens; - std::vector mNextDraftLens; - DraftPaths mNextDraftPaths; - - DraftTokensVec mLastDraftTokens; - std::vector mLastDraftLens; - DraftPaths mLastDraftPaths; - - std::vector mAcceptedTokens; - std::vector mAcceptedLens; - std::vector mAcceptedPathIds; - - std::vector>> mMasks; -}; - -template -class EagleDecodingLayerTest : public testing::Test -{ -private: - void SetUp() override; - -private: - SamplingParams mSamplingParams; - - // Outputs - TensorPtr mOutputIds; - TensorPtr mSeqLengths; - TensorPtr mOutputNextDraftTokens; - TensorPtr mOutputUnpackedNextDraftTokens; - TensorPtr mAcceptedLengths; - TensorPtr mNextPosIds; - TensorPtr mPrevDraftLengths; - TensorPtr mNextDraftLengths; - TensorPtr mNextGenerationLengths; - TensorPtr mNextGenerationLengthsHost; - TensorPtr mAcceptedLengthCumSum; - TensorPtr mPathsOffsets; - TensorPtr mPackedMasks; - TensorPtr mRandomDataSample; - TensorPtr mRandomDataValidation; - TensorPtr mOutputTemperatures; - TensorPtr mOutputNextDraftPaths; - TensorPtr mEagleNetCtxRequestTypesHost; - TensorPtr mEagleNetCtxContextLengthsHost; - TensorPtr mEagleNetCtxPastKeyValueLengthsHost; - TensorPtr mEagleNetGenRequestTypesHost; - TensorPtr mEagleNetGenContextLengthsHost; - TensorPtr mEagleNetGenPastKeyValueLengthsHost; - - // inputs - TensorPtr mBatchSlots; - TensorPtr mEndIds; - - TensorPtr mInputNextDraftTokens; - TensorPtr mInputNextDraftLens; - TensorPtr mInputNextDraftPaths; - TensorPtr mInputLastDraftTokens; - TensorPtr mInputLastDraftLens; - TensorPtr mInputLastDraftPaths; - TensorPtr mInputAcceptedTokens; - TensorPtr mInputAcceptedLens; - TensorPtr mInputAcceptedPathIds; - TensorPtr mChunkedContextNextTokens; - - // Setup params - std::vector mRandomSeeds; - std::vector mTemperatures; - - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - std::shared_ptr> mEagleLayer; - std::shared_ptr mDecodingWorkspace; - - EagleDummyNetwork mNetwork; - -private: - void allocateBuffers(); - - void setup(); - - std::shared_ptr createInputTensors(); - - std::shared_ptr createOutputTensors(); - - void checkLayerResult(); - -public: - void runTest(std::vector const& prompts, std::vector const& predictions, - std::vector const& nextDraftLetters, std::vector const& lastDraftLetters, - SamplingParams& params); -}; - -typedef testing::Types FloatAndHalfTypes; - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp b/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp deleted file mode 100644 index 04d05e0d16a9..000000000000 --- a/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp +++ /dev/null @@ -1,1661 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tests/unit_tests/layers/explicitDraftTokensLayerTest.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/speculativeDecodingModule.h" - -#include "tensorrt_llm/common/tllmDataType.h" - -#include -#include -#include - -namespace tensorrt_llm::tests::layers -{ -// TODO verify context + gen mix - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using namespace tensorrt_llm::common; - -namespace tk = tensorrt_llm::kernels; -namespace tksd = tensorrt_llm::kernels::speculative_decoding; -namespace trk = tensorrt_llm::runtime::kernels; - -TokensVec ExplicitDraftTokensDummyNetwork::tokenize(std::string const& letters) const -{ - TokensVec tokens; - for (char c : letters) - { - tokens.push_back(static_cast(c)); - } - return tokens; -} - -std::string ExplicitDraftTokensDummyNetwork::detokenize(TokensVec const& tokens) const -{ - std::string letters; - for (int token : tokens) - { - letters += static_cast(token); - } - return letters; -} - -DraftTokensVec ExplicitDraftTokensDummyNetwork::draftLettersToTokens(DraftLettersVec const& draftLetters) const -{ - DraftTokensVec draftTokens(draftLetters.size()); - for (SizeType32 bi = 0; bi < draftLetters.size(); ++bi) - { - draftTokens[bi].resize(draftLetters[bi].size()); - for (SizeType32 pi = 0; pi < draftLetters[bi].size(); ++pi) - { - draftTokens[bi][pi] = tokenize(draftLetters[bi][pi]); - } - } - return draftTokens; -} - -SizeType32 ExplicitDraftTokensDummyNetwork::longestCommonPrefixLength(TokensVec const& a, TokensVec const& b) const -{ - SizeType32 minLength = std::min(a.size(), b.size()); - SizeType32 idx = 0; - while (idx < minLength && a[idx] == b[idx]) - { - ++idx; - } - return idx; -} - -SizeType32 ExplicitDraftTokensDummyNetwork::computeCompressedVectorAndIndices(TokensVec& compressedVector, - std::vector& packedPosIds, DraftTokensIndices& indices, std::vector const& vectors, - SizeType32 basePosId) -{ - TokensVec localCompressedVector; - std::vector localPackedPosIds; - std::vector> localIndices; - - // FIXME always take the 1st beam as the reference. Is that correct? - // Add whole first vector to compressed vector - localCompressedVector = vectors[0]; - // All indices of first vector. - localIndices.push_back(std::vector(localCompressedVector.size())); - for (SizeType32 ti = 0; ti < localCompressedVector.size(); ++ti) - { - localIndices[0][ti] = ti; - // Set local to batch packed pos ids. - localPackedPosIds.push_back(basePosId + ti); - } - - // Starting from the 1st path. - for (SizeType32 pi = 1; pi < vectors.size(); ++pi) - { - // Match path to compressed vector (aka path 0). - auto const prefixLength = longestCommonPrefixLength(localCompressedVector, vectors[pi]); - localIndices.push_back(std::vector(vectors[pi].size())); - // Set indices of the matched prefix. - for (SizeType32 ti = 0; ti < prefixLength; ++ti) - { - localIndices[pi][ti] = ti; - } - // For non-matched part. - for (SizeType32 ti = prefixLength; ti < vectors[pi].size(); ++ti) - { - // Add new tokens to compressed vector. - localCompressedVector.push_back(vectors[pi][ti]); - // Set new pos ids. - localPackedPosIds.push_back(basePosId + ti); - // Set their indices. - localIndices[pi][ti] = localCompressedVector.size() - 1; - } - } - - compressedVector.insert(compressedVector.end(), localCompressedVector.begin(), localCompressedVector.end()); - packedPosIds.insert(packedPosIds.end(), localPackedPosIds.begin(), localPackedPosIds.end()); - indices.push_back(localIndices); - return static_cast(localCompressedVector.size()); -} - -void ExplicitDraftTokensDummyNetwork::createNextMasks( - DraftTokensIndices const& indices, DraftTokensVec const& draftTokens, SizeType32 maxGenLength) -{ - for (SizeType32 bi = 0; bi < indices.size(); ++bi) - { - std::vector> localMask(maxGenLength, std::vector(maxGenLength)); - // Create fill diagonal. - for (SizeType32 ti = 0; ti < maxGenLength; ++ti) - { - localMask[ti][ti] = true; - } - - SizeType32 rowIdx = 0; - for (SizeType32 pi = 0; pi < draftTokens[bi].size(); ++pi) - { - auto const prefixLength = pi == 0 ? 0 : longestCommonPrefixLength(draftTokens[bi][0], draftTokens[bi][pi]); - for (SizeType32 ti = 0; ti < draftTokens[bi][pi].size(); ++ti) - { - auto const index = indices[bi][pi][ti]; - // If we are in the "prefix" part of the sequence skip it as it does not represent real mask row. - if (ti < prefixLength) - { - continue; - } - // Fill lower triangular part according to the prefix. - for (SizeType32 tti = 0; tti < ti; ++tti) - { - localMask[rowIdx][indices[bi][pi][tti]] = true; - } - rowIdx++; - } - } - mMasks.push_back(localMask); - } -} - -void ExplicitDraftTokensDummyNetwork::compressTokens(TokensVec& compressedVector, std::vector& packedPosIds, - DraftTokensIndices& indices, std::vector& generationLengths, DraftTokensVec const& draftTokens, - std::vector const& basePosIds) -{ - generationLengths.resize(draftTokens.size()); - for (SizeType32 bi = 0; bi < draftTokens.size(); ++bi) - { - auto numGeneratedTokens = computeCompressedVectorAndIndices( - compressedVector, packedPosIds, indices, draftTokens[bi], basePosIds[bi]); - generationLengths[bi] = numGeneratedTokens; - } - // Pad vectors to the maximum size - auto const padSize - = mSamplingParams.getMaxDecodingTokens() * mSamplingParams.getBatchSize() - compressedVector.size(); - compressedVector.insert(compressedVector.end(), padSize, mSamplingParams.getPadId()); - packedPosIds.insert(packedPosIds.end(), padSize, 0); -} - -void ExplicitDraftTokensDummyNetwork::acceptTokens(std::vector const& predictionTokens, - DraftTokensVec const& lastDraftTokens, DraftTokensVec const& nextDraftTokens) -{ - TLLM_CHECK_WITH_INFO(predictionTokens.size() == lastDraftTokens.size(), - "Batch size of predictions (%d) does not match the batch size of last draft tokens (%d)", - static_cast(predictionTokens.size()), static_cast(lastDraftTokens.size())); - TLLM_CHECK_WITH_INFO(predictionTokens.size() == nextDraftTokens.size(), - "Batch size of predictions (%d) does not match the batch size of next draft tokens (%d)", - static_cast(predictionTokens.size()), static_cast(nextDraftTokens.size())); - mBestPathLengths.resize(predictionTokens.size()); - mBestPathIndices.resize(predictionTokens.size()); - // Needed for unit test of ExplicitDraftTokensDummyNetwork only. - if (mOutputIds.size() == 0) - { - mOutputIds.resize(lastDraftTokens.size()); - } - for (SizeType32 bi = 0; bi < predictionTokens.size(); ++bi) - { - SizeType32 maxMatchLen = -1; - SizeType32 maxMatchIdx = -1; - // Find path with largest prefix shared with the predicted tokens. - for (SizeType32 pi = 0; pi < lastDraftTokens[bi].size(); ++pi) - { - TLLM_CHECK_WITH_INFO(predictionTokens[bi][0] == lastDraftTokens[bi][pi][0], - "First token of prediction and draft token must match"); - auto const matchLen = longestCommonPrefixLength(lastDraftTokens[bi][pi], predictionTokens[bi]); - if (matchLen > maxMatchLen) - { - maxMatchLen = matchLen; - maxMatchIdx = pi; - } - } - mBestPathLengths[bi] = maxMatchLen; - mBestPathIndices[bi] = maxMatchIdx; - // Update output ids. First draft token is already counted in outputs - mOutputIds[bi].insert(mOutputIds[bi].end(), lastDraftTokens[bi][maxMatchIdx].begin() + 1, - lastDraftTokens[bi][maxMatchIdx].begin() + maxMatchLen); - mOutputIds[bi].push_back(nextDraftTokens[bi][0][0]); - } -} - -void ExplicitDraftTokensDummyNetwork::forward(SamplingParams const& params, - std::vector const& promptsLetters, std::vector const& predictionLetters, - DraftLettersVec const& nextDraftLetters, DraftLettersVec const& lastDraftLetters) -{ - mSamplingParams = params; - - TLLM_CHECK(params.getBatchSize() == promptsLetters.size()); - TLLM_CHECK(params.getBatchSize() == predictionLetters.size()); - TLLM_CHECK(params.getBatchSize() == nextDraftLetters.size()); - TLLM_CHECK(params.getBatchSize() == lastDraftLetters.size()); - - // Tokenize - mNextDraftTokens = draftLettersToTokens(nextDraftLetters); - mLastDraftTokens = draftLettersToTokens(lastDraftLetters); - std::vector predictionTokens; - for (SizeType32 bi = 0; bi < predictionLetters.size(); ++bi) - { - predictionTokens.push_back(tokenize(predictionLetters[bi])); - mPrompts.push_back(tokenize(promptsLetters[bi])); - } - - std::vector basePosIds; - for (auto const& prompt : mPrompts) - { - basePosIds.push_back(prompt.size()); - } - - mOutputIds = mPrompts; - - // Make compressed tensors and pos ids for the current and next tokens - compressTokens(mNextCompressedVector, mNextPackedPosIds, mNextDraftTokenIndices, mNextGenerationLengths, - mNextDraftTokens, basePosIds); - compressTokens(mLastCompressedVector, mLastPackedPosIds, mLastDraftTokenIndices, mLastGenerationLengths, - mLastDraftTokens, basePosIds); - - mMaxNextGenLength = *std::max_element(mNextGenerationLengths.begin(), mNextGenerationLengths.end()); - - acceptTokens(predictionTokens, mLastDraftTokens, mNextDraftTokens); - - createNextMasks(mNextDraftTokenIndices, mNextDraftTokens, mMaxNextGenLength); -} - -TEST(ExplicitDraftTokensDummyNetworkTest, tokenizeTest) -{ - ExplicitDraftTokensDummyNetwork network; - - { - auto tokens = network.tokenize("hello world"); - EXPECT_EQ(tokens, std::vector({104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100})); - } - { - DraftLettersVec lettersVec = {{"hello world", "hello"}, {"world"}}; - auto draftTokens = network.draftLettersToTokens(lettersVec); - ASSERT_EQ(draftTokens.size(), 2); - ASSERT_EQ(draftTokens[0].size(), 2); - ASSERT_EQ(draftTokens[1].size(), 1); - EXPECT_EQ(draftTokens[0][0], std::vector({104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100})); - EXPECT_EQ(draftTokens[0][1], std::vector({104, 101, 108, 108, 111})); - EXPECT_EQ(draftTokens[1][0], std::vector({119, 111, 114, 108, 100})); - } -} - -TEST(ExplicitDraftTokensDummyNetworkTest, detokenizeTest) -{ - ExplicitDraftTokensDummyNetwork network; - - { - auto letters - = network.detokenize(std::vector({104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100})); - EXPECT_EQ(letters, "hello world"); - } -} - -TEST(ExplicitDraftTokensDummyNetworkTest, longestCommonPrefixLengthTest) -{ - ExplicitDraftTokensDummyNetwork network; - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {1, 2}), 2); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {1, 2, 3}), 3); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {1, 5, 6}), 1); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {2, 5, 6}), 0); - EXPECT_EQ(network.longestCommonPrefixLength({1, 2, 3}, {}), 0); -} - -TEST(ExplicitDraftTokensDummyNetworkTest, computeCompressedVectorAndIndicesTest) -{ - ExplicitDraftTokensDummyNetwork network; - - { - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - - SizeType32 basePosId{0}; - - std::vector> tokens = {{0, 1, 2, 3}}; - - auto const totalGen - = network.computeCompressedVectorAndIndices(compressedVector, packedPosIds, indices, tokens, basePosId); - - EXPECT_EQ(totalGen, 4); - EXPECT_EQ(compressedVector, std::vector({0, 1, 2, 3})); - EXPECT_EQ(packedPosIds, std::vector({0, 1, 2, 3})); - ASSERT_EQ(indices.size(), 1); - ASSERT_EQ(indices[0].size(), 1); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - } - - { - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - - SizeType32 basePosId{0}; - - std::vector> tokens = {{0, 1, 2, 3}, {0, 2, 3, 4}}; - - auto const totalGen - = network.computeCompressedVectorAndIndices(compressedVector, packedPosIds, indices, tokens, basePosId); - - EXPECT_EQ(totalGen, 7); - EXPECT_EQ(compressedVector, std::vector({0, 1, 2, 3, 2, 3, 4})); - EXPECT_EQ(packedPosIds, std::vector({0, 1, 2, 3, 1, 2, 3})); - ASSERT_EQ(indices.size(), 1); - ASSERT_EQ(indices[0].size(), 2); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - EXPECT_EQ(indices[0][1], std::vector({0, 4, 5, 6})); - } - - { - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - - SizeType32 basePosId{0}; - - std::vector> tokens = {{0, 1, 2, 3}, {0, 1, 6, 2}, {0, 5, 6, 2}}; - - auto const totalGen - = network.computeCompressedVectorAndIndices(compressedVector, packedPosIds, indices, tokens, basePosId); - - EXPECT_EQ(totalGen, 9); - EXPECT_EQ(compressedVector, std::vector({0, 1, 2, 3, 6, 2, 5, 6, 2})); - EXPECT_EQ(packedPosIds, std::vector({0, 1, 2, 3, 2, 3, 1, 2, 3})); - ASSERT_EQ(indices.size(), 1); - ASSERT_EQ(indices[0].size(), 3); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - EXPECT_EQ(indices[0][1], std::vector({0, 1, 4, 5})); - EXPECT_EQ(indices[0][2], std::vector({0, 6, 7, 8})); - } - - { - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - - SizeType32 basePosId{10}; - - std::vector> tokens = {{0, 1, 2, 3}, {0, 1, 6, 2}, {0, 5, 6, 2}}; - - auto const totalGen - = network.computeCompressedVectorAndIndices(compressedVector, packedPosIds, indices, tokens, basePosId); - - EXPECT_EQ(totalGen, 9); - EXPECT_EQ(compressedVector, std::vector({0, 1, 2, 3, 6, 2, 5, 6, 2})); - EXPECT_EQ(packedPosIds, std::vector({10, 11, 12, 13, 12, 13, 11, 12, 13})); - ASSERT_EQ(indices.size(), 1); - ASSERT_EQ(indices[0].size(), 3); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - EXPECT_EQ(indices[0][1], std::vector({0, 1, 4, 5})); - EXPECT_EQ(indices[0][2], std::vector({0, 6, 7, 8})); - } -} - -TEST(ExplicitDraftTokensDummyNetworkTest, compressTokensTest) -{ - { - ExplicitDraftTokensDummyNetwork network; - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - std::vector genLengths; - - SamplingParams params; - params.setBatchSize(1); - params.setMaxNumPaths(1); - params.setMaxDraftPathLen(6); - network.setSamplingParams(params); - - DraftTokensVec tokens = {{{0, 1, 2, 3}}}; - - std::vector basePosIds = {0}; - - network.compressTokens(compressedVector, packedPosIds, indices, genLengths, tokens, basePosIds); - - EXPECT_EQ(compressedVector, std::vector({0, 1, 2, 3, -1, -1, -1})); - EXPECT_EQ(packedPosIds, std::vector({0, 1, 2, 3, 0, 0, 0})); - ASSERT_EQ(indices.size(), 1); - ASSERT_EQ(indices[0].size(), 1); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - ASSERT_EQ(genLengths.size(), 1); - EXPECT_EQ(genLengths[0], 4); - - network.createNextMasks(indices, tokens, 4); - auto masks = network.getNextMasks(); - ASSERT_EQ(masks.size(), 1); - ASSERT_EQ(masks[0].size(), 4); - ASSERT_EQ(masks[0][0].size(), 4); - - EXPECT_EQ(masks[0][0], std::vector({true, false, false, false})); - EXPECT_EQ(masks[0][1], std::vector({true, true, false, false})); - EXPECT_EQ(masks[0][2], std::vector({true, true, true, false})); - EXPECT_EQ(masks[0][3], std::vector({true, true, true, true})); - } - - { - ExplicitDraftTokensDummyNetwork network; - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - std::vector genLengths; - - SamplingParams params; - params.setBatchSize(2); - params.setMaxNumPaths(1); - params.setMaxDraftPathLen(6); - network.setSamplingParams(params); - - std::vector basePosIds = {10, 10}; - - DraftTokensVec tokens = {{{0, 1, 2, 3}}, {{0, 1, 2, 3}}}; - - network.compressTokens(compressedVector, packedPosIds, indices, genLengths, tokens, basePosIds); - - EXPECT_EQ(compressedVector, std::vector({0, 1, 2, 3, 0, 1, 2, 3, -1, -1, -1, -1, -1, -1})); - EXPECT_EQ(packedPosIds, std::vector({10, 11, 12, 13, 10, 11, 12, 13, 0, 0, 0, 0, 0, 0})); - ASSERT_EQ(indices.size(), 2); - ASSERT_EQ(indices[0].size(), 1); - ASSERT_EQ(indices[1].size(), 1); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - EXPECT_EQ(indices[1][0], std::vector({0, 1, 2, 3})); - ASSERT_EQ(genLengths.size(), 2); - EXPECT_EQ(genLengths[0], 4); - EXPECT_EQ(genLengths[1], 4); - - network.createNextMasks(indices, tokens, 4); - auto masks = network.getNextMasks(); - ASSERT_EQ(masks.size(), 2); - ASSERT_EQ(masks[0].size(), 4); - ASSERT_EQ(masks[1].size(), 4); - ASSERT_EQ(masks[0][0].size(), 4); - ASSERT_EQ(masks[1][0].size(), 4); - - EXPECT_EQ(masks[0][0], std::vector({true, false, false, false})); - EXPECT_EQ(masks[0][1], std::vector({true, true, false, false})); - EXPECT_EQ(masks[0][2], std::vector({true, true, true, false})); - EXPECT_EQ(masks[0][3], std::vector({true, true, true, true})); - - EXPECT_EQ(masks[1][0], std::vector({true, false, false, false})); - EXPECT_EQ(masks[1][1], std::vector({true, true, false, false})); - EXPECT_EQ(masks[1][2], std::vector({true, true, true, false})); - EXPECT_EQ(masks[1][3], std::vector({true, true, true, true})); - } - { - ExplicitDraftTokensDummyNetwork network; - std::vector compressedVector; - std::vector packedPosIds; - DraftTokensIndices indices; - std::vector genLengths; - - SamplingParams params; - params.setBatchSize(2); - params.setMaxNumPaths(3); - params.setMaxDraftPathLen(4); - network.setSamplingParams(params); - - std::vector basePosIds = {10, 0}; - - DraftTokensVec tokens = {{{0, 1, 2, 3}, {0, 1, 6, 2}, {0, 5, 6, 2}}, {{0, 1, 2, 3}, {0, 1, 2, 4}}}; - - network.compressTokens(compressedVector, packedPosIds, indices, genLengths, tokens, basePosIds); - - EXPECT_EQ(compressedVector, - std::vector( - {0, 1, 2, 3, 6, 2, 5, 6, 2, 0, 1, 2, 3, 4, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1})); - EXPECT_EQ(packedPosIds, - std::vector( - {10, 11, 12, 13, 12, 13, 11, 12, 13, 0, 1, 2, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})); - ASSERT_EQ(indices.size(), 2); - ASSERT_EQ(indices[0].size(), 3); - ASSERT_EQ(indices[1].size(), 2); - EXPECT_EQ(indices[0][0], std::vector({0, 1, 2, 3})); - EXPECT_EQ(indices[0][1], std::vector({0, 1, 4, 5})); - EXPECT_EQ(indices[0][2], std::vector({0, 6, 7, 8})); - EXPECT_EQ(indices[1][0], std::vector({0, 1, 2, 3})); - EXPECT_EQ(indices[1][1], std::vector({0, 1, 2, 4})); - ASSERT_EQ(genLengths.size(), 2); - EXPECT_EQ(genLengths[0], 9); - EXPECT_EQ(genLengths[1], 5); - - network.createNextMasks(indices, tokens, 9); - auto masks = network.getNextMasks(); - ASSERT_EQ(masks.size(), 2); - ASSERT_EQ(masks[0].size(), 9); - ASSERT_EQ(masks[1].size(), 9); - ASSERT_EQ(masks[0][0].size(), 9); - ASSERT_EQ(masks[1][0].size(), 9); - - EXPECT_EQ(masks[0][0], std::vector({true, false, false, false, false, false, false, false, false})); - EXPECT_EQ(masks[0][1], std::vector({true, true, false, false, false, false, false, false, false})); - EXPECT_EQ(masks[0][2], std::vector({true, true, true, false, false, false, false, false, false})); - EXPECT_EQ(masks[0][3], std::vector({true, true, true, true, false, false, false, false, false})); - EXPECT_EQ(masks[0][4], std::vector({true, true, false, false, true, false, false, false, false})); - EXPECT_EQ(masks[0][5], std::vector({true, true, false, false, true, true, false, false, false})); - EXPECT_EQ(masks[0][6], std::vector({true, false, false, false, false, false, true, false, false})); - EXPECT_EQ(masks[0][7], std::vector({true, false, false, false, false, false, true, true, false})); - EXPECT_EQ(masks[0][8], std::vector({true, false, false, false, false, false, true, true, true})); - - EXPECT_EQ(masks[1][0], std::vector({true, false, false, false, false, false, false, false, false})); - EXPECT_EQ(masks[1][1], std::vector({true, true, false, false, false, false, false, false, false})); - EXPECT_EQ(masks[1][2], std::vector({true, true, true, false, false, false, false, false, false})); - EXPECT_EQ(masks[1][3], std::vector({true, true, true, true, false, false, false, false, false})); - EXPECT_EQ(masks[1][4], std::vector({true, true, true, false, true, false, false, false, false})); - EXPECT_EQ(masks[1][5], std::vector({false, false, false, false, false, true, false, false, false})); - EXPECT_EQ(masks[1][6], std::vector({false, false, false, false, false, false, true, false, false})); - EXPECT_EQ(masks[1][7], std::vector({false, false, false, false, false, false, false, true, false})); - EXPECT_EQ(masks[1][8], std::vector({false, false, false, false, false, false, false, false, true})); - } -} - -TEST(ExplicitDraftTokensDummyNetworkTest, acceptTokensTest) -{ - { - ExplicitDraftTokensDummyNetwork network; - std::vector predictionTokens = {network.tokenize("how things")}; - DraftLettersVec lastDraftLetters = {{"how do ", "how are", "however", "hello w"}}; - DraftLettersVec nextDraftLetters = {{"things ", "that is", "to crea", "touchab"}}; - auto lastDraftTokens = network.draftLettersToTokens(lastDraftLetters); - auto nextDraftTokens = network.draftLettersToTokens(nextDraftLetters); - - network.acceptTokens(predictionTokens, lastDraftTokens, nextDraftTokens); - - auto bestPathLengths = network.getBestPathLengths(); - auto bestPathIndices = network.getBestPathIndices(); - auto outputIds = network.getOutputIds(); - - ASSERT_EQ(bestPathLengths.size(), 1); - ASSERT_EQ(bestPathIndices.size(), 1); - ASSERT_EQ(outputIds.size(), 1); - EXPECT_EQ(bestPathLengths[0], 4); - EXPECT_EQ(bestPathIndices[0], 0); - EXPECT_EQ(network.detokenize(outputIds[0]), "ow t"); - } - - { - ExplicitDraftTokensDummyNetwork network; - std::vector predictionTokens = {network.tokenize("however you")}; - DraftLettersVec lastDraftLetters = {{"how do ", "how tho", "however", "hello w"}}; - DraftLettersVec nextDraftLetters = {{" increme", " introdu", " i = 0; ", " importa"}}; - auto lastDraftTokens = network.draftLettersToTokens(lastDraftLetters); - auto nextDraftTokens = network.draftLettersToTokens(nextDraftLetters); - - network.acceptTokens(predictionTokens, lastDraftTokens, nextDraftTokens); - - auto bestPathLengths = network.getBestPathLengths(); - auto bestPathIndices = network.getBestPathIndices(); - auto outputIds = network.getOutputIds(); - - ASSERT_EQ(bestPathLengths.size(), 1); - ASSERT_EQ(bestPathIndices.size(), 1); - ASSERT_EQ(outputIds.size(), 1); - EXPECT_EQ(bestPathLengths[0], 7); - EXPECT_EQ(bestPathIndices[0], 2); - EXPECT_EQ(network.detokenize(outputIds[0]), "owever "); - } - - { - ExplicitDraftTokensDummyNetwork network; - std::vector predictionTokens = {network.tokenize("how things")}; - DraftLettersVec lastDraftLetters = {{"heruist", "habit i", "handove", "hammer "}}; - DraftLettersVec nextDraftLetters = {{"oatmeal", "ocean b", "occupat", "oblivio"}}; - auto lastDraftTokens = network.draftLettersToTokens(lastDraftLetters); - auto nextDraftTokens = network.draftLettersToTokens(nextDraftLetters); - - network.acceptTokens(predictionTokens, lastDraftTokens, nextDraftTokens); - - auto bestPathLengths = network.getBestPathLengths(); - auto bestPathIndices = network.getBestPathIndices(); - auto outputIds = network.getOutputIds(); - - ASSERT_EQ(bestPathLengths.size(), 1); - ASSERT_EQ(bestPathIndices.size(), 1); - ASSERT_EQ(outputIds.size(), 1); - EXPECT_EQ(bestPathLengths[0], 1); - EXPECT_EQ(bestPathIndices[0], 0); - EXPECT_EQ(network.detokenize(outputIds[0]), "o"); - } -} - -///////////////////////////////////////////////////////////////////////////////////////////////////////////////// - -template -void ExplicitDraftTokensLayerTest::SetUp() -{ - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); -} - -template -void ExplicitDraftTokensLayerTest::allocateBuffers() -{ - using DataType = typename T::DataType; - auto const dataType = TRTDataType::value; - - auto speculativeDecodingModule = std::make_shared(mSamplingParams.getMaxDraftPathLen(), - mSamplingParams.getMaxDecodingDraftTokens(), mSamplingParams.getMaxNumPaths()); - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain(mSamplingParams.getMaxBatchSize(), 1, - mSamplingParams.getVocabSize(), mSamplingParams.getVocabSize(), speculativeDecodingModule); - - mExplicitDraftTokensLayer = std::make_shared>( - decodingDomain, mBufferManager); - - // outputs - mOutputIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxSeqLen()}), - tensorrt_llm::DataType::kINT32); - - mSeqLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mAcceptedLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mPrevDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mAcceptedLengthCumSum = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() + 1}), tensorrt_llm::DataType::kINT32); - - mOutputNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - tensorrt_llm::DataType::kINT32); - - mOutputPositionIdsBase = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mRandomDataSample = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), dataType); - - mRandomDataValidation - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize(), - mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxDraftPathLen()}), - dataType); - - mPackedMasks = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), - static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))}), - tensorrt_llm::DataType::kINT32); - - mNextPosIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - - mOutputUnpackedNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mOutputUnpackedNextDraftIndices = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mOutputDraftProbs = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), - mSamplingParams.getMaxDraftPathLen(), mSamplingParams.getVocabSize()}), - dataType); - - mOutputTemperatures = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), dataType); - - mOutputGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mOutputGenerationLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mMaxGenLengthHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - // inputs - mBatchSlots = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - mTokensPerStep = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mPathsOffsets = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() * mSamplingParams.getMaxDraftPathLen()}), - tensorrt_llm::DataType::kINT32); - - mMasks = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), - mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kBOOL); - - mInputNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mLastDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mPackedPosIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - - mBestPathLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mBestPathIndices = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mSpecDecodingGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextFlatTokens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() * mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - - mInputPositionIdsBase = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mNextDraftIndices = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mLastDraftIndices = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - - mNextDraftProbs = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), - mSamplingParams.getMaxDraftPathLen(), mSamplingParams.getVocabSize()}), - dataType); - - mEndIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - - mMaxGenLengthDevice = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - - // Packed inputs - mMaxGenerationLength = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - mCumSumGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - - // Packed outputs - mPackedPositionIdsBase = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - mPackedGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - mPackedRandomDataSample = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), dataType); - mPackedRandomDataVerification = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxDraftPathLen()}), - dataType); - mPackedNextDraftTokens = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - mPackedNextDraftIndices = BufferManager::pinnedPool( - ITensor::makeShape( - {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - tensorrt_llm::DataType::kINT32); - mPackedPackedMasks = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), - static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))}), - tensorrt_llm::DataType::kINT32); - mPackedPositionOffsets = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - mPackedPackedPosIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - tensorrt_llm::DataType::kINT32); - mPackedDraftProbs = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), - mSamplingParams.getMaxDraftPathLen(), mSamplingParams.getVocabSize()}), - dataType); - mPackedTemperatures = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), dataType); - mDecodingWorkspace = std::make_shared(mBufferManager, decodingDomain, - TRTDataType::value, mExplicitDraftTokensLayer->getWorkspaceSize()); -} - -template -void ExplicitDraftTokensLayerTest::setup() -{ - using DataType = typename T::DataType; - // outputs - trk::invokeFill(*mOutputIds, TokenIdType{-1}, *mStream); - trk::invokeFill(*mSeqLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mAcceptedLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mAcceptedLengthCumSum, SizeType32{-1}, *mStream); - trk::invokeFill(*mOutputNextDraftTokens, TokenIdType{-1}, *mStream); - trk::invokeFill(*mOutputPositionIdsBase, SizeType32{0}, *mStream); - trk::invokeFill(*mRandomDataSample, DataType{0}, *mStream); - trk::invokeFill(*mRandomDataValidation, DataType{0}, *mStream); - trk::invokeFill(*mPackedMasks, SizeType32{0}, *mStream); - trk::invokeFill(*mNextPosIds, SizeType32{0}, *mStream); - trk::invokeFill(*mOutputUnpackedNextDraftTokens, TokenIdType{-1}, *mStream); - trk::invokeFill(*mOutputUnpackedNextDraftIndices, SizeType32{0}, *mStream); - trk::invokeFill(*mEndIds, TokenIdType{-1}, *mStream); - - auto inDraftProbs = BufferRange(*mNextDraftProbs); - - std::mt19937 gen(42); - std::uniform_real_distribution distr(0.0, 1.0); - std::generate( - inDraftProbs.begin(), inDraftProbs.end(), [&gen, &distr]() { return static_cast(distr(gen)); }); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - batchSlotsPtr[bi] = 2 * bi; - } - - auto setupParams = std::make_shared(); - mRandomSeeds = std::vector(mSamplingParams.getBatchSize()); - mTemperatures = std::vector(mSamplingParams.getBatchSize()); - - std::mt19937 generator(42); - std::uniform_int_distribution seedDistr(1, 1000); - std::uniform_real_distribution temperatureDistr(0.001f, 1.f); - std::generate( - mRandomSeeds.begin(), mRandomSeeds.end(), [&generator, &seedDistr]() { return seedDistr(generator); }); - std::generate(mTemperatures.begin(), mTemperatures.end(), - [&generator, &temperatureDistr]() { return temperatureDistr(generator); }); - setupParams->randomSeed = mRandomSeeds; - setupParams->temperature = mTemperatures; - setupParams->randomDataSample = mRandomDataSample; - setupParams->temperatures = mOutputTemperatures; - setupParams->dtype = TRTDataType::value; - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mExplicitDraftTokensLayer->setup(mSamplingParams.getBatchSize(), 1, mBatchSlots, setupParams, mDecodingWorkspace); - - mStream->synchronize(); - - mBestPathLengths = mBufferManager->copyFrom(mNetwork.getBestPathLengths(), - ITensor::makeShape({mSamplingParams.getBatchSize()}), runtime::MemoryType::kPINNEDPOOL); - mBestPathIndices = mBufferManager->copyFrom(mNetwork.getBestPathIndices(), - ITensor::makeShape({mSamplingParams.getBatchSize()}), runtime::MemoryType::kPINNEDPOOL); - mPackedPosIds = mBufferManager->copyFrom(mNetwork.getNextPackedPosId(), - ITensor::makeShape({mSamplingParams.getMaxDecodingTokens() * mSamplingParams.getBatchSize()}), - runtime::MemoryType::kPINNEDPOOL); - - auto const nextDraftTokens = mNetwork.getNextDraftTokens(); - auto const lastDraftTokens = mNetwork.getLastDraftTokens(); - auto const nextDraftIndices = mNetwork.getNextDraftIndices(); - auto const lastDraftIndices = mNetwork.getLastDraftIndices(); - auto sequenceLength = BufferRange(*mSeqLengths); - auto nextDraftTokensRange = BufferRange(*mInputNextDraftTokens); - auto lastDraftTokensRange = BufferRange(*mLastDraftTokens); - auto nextDraftIndicesRange = BufferRange(*mNextDraftIndices); - auto lastDraftIndicesRange = BufferRange(*mLastDraftIndices); - auto inputPositionIdsBase = BufferRange(*mInputPositionIdsBase); - - auto outputIds = BufferRange(*mOutputIds); - auto generationLengths = mNetwork.getNextGenerationLengths(); - auto prompts = mNetwork.getPrompts(); - for (SizeType32 bi = 0; bi < nextDraftTokens.size(); ++bi) - { - for (SizeType32 pi = 0; pi < nextDraftTokens[bi].size(); ++pi) - { - for (SizeType32 ti = 0; ti < nextDraftTokens[bi][pi].size(); ++ti) - { - auto idx = flat_index3(bi, pi, ti, mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()); - nextDraftTokensRange[idx] = nextDraftTokens[bi][pi][ti]; - lastDraftTokensRange[idx] = lastDraftTokens[bi][pi][ti]; - nextDraftIndicesRange[idx] = nextDraftIndices[bi][pi][ti]; - lastDraftIndicesRange[idx] = lastDraftIndices[bi][pi][ti]; - } - } - bufferCast(*mSpecDecodingGenerationLengths)[bi] = generationLengths[bi]; - - sequenceLength[batchSlotsPtr[bi]] = prompts[bi].size(); - std::copy(prompts[bi].begin(), prompts[bi].end(), - outputIds.begin() + batchSlotsPtr[bi] * mSamplingParams.getMaxSeqLen()); - - inputPositionIdsBase[bi] = prompts[bi].size(); - } - - auto nextFlatTokens = mNetwork.getNextFlatTokens(); - TLLM_LOG_DEBUG("Next flat tokens are \"%s\"", mNetwork.detokenize(nextFlatTokens).c_str()); - auto nextFlatTokensRange = BufferRange(*mNextFlatTokens); - std::copy(nextFlatTokens.begin(), nextFlatTokens.end(), nextFlatTokensRange.begin()); - - auto const masks = mNetwork.getNextMasks(); - auto masksRange = BufferRange(*mMasks); - auto const maxGenLength = mNetwork.getMaxNextGenerationLength(); - bufferCast(*mMaxGenerationLength)[0] = maxGenLength; - for (SizeType32 bi = 0; bi < masks.size(); ++bi) - { - TLLM_CHECK(maxGenLength == masks[bi].size()); - for (SizeType32 ri = 0; ri < masks[bi].size(); ++ri) - { - TLLM_CHECK(maxGenLength == masks[bi][ri].size()); - for (SizeType32 ci = 0; ci < masks[bi][ri].size(); ++ci) - { - masksRange[bi * maxGenLength * maxGenLength + ri * maxGenLength + ci] = masks[bi][ri][ci]; - } - } - } -} - -template -std::shared_ptr ExplicitDraftTokensLayerTest::createInputTensors() -{ - auto forwardParams - = std::make_shared(mEndIds, mBatchSlots, mSamplingParams.getBatchSize()); - - forwardParams->seqSlots = mBatchSlots; - - forwardParams->masks = mMasks; - - forwardParams->nextDraftTokens = mInputNextDraftTokens; - - forwardParams->nextDraftIndices = mNextDraftIndices; - - forwardParams->lastDraftTokens = mLastDraftTokens; - - forwardParams->lastDraftIndices = mLastDraftIndices; - - forwardParams->packedPosIds = mPackedPosIds; - - forwardParams->bestPathLengths = mBestPathLengths; - - forwardParams->bestPathIndices = mBestPathIndices; - - forwardParams->generationLengths = mSpecDecodingGenerationLengths; - - forwardParams->nextFlatTokens = mNextFlatTokens; - - forwardParams->positionIdsBase = mInputPositionIdsBase; - - forwardParams->nextDraftProbs = mNextDraftProbs; - - forwardParams->maxGenLengthDevice = mMaxGenLengthDevice; - - return forwardParams; -} - -template -std::shared_ptr ExplicitDraftTokensLayerTest::createOutputTensors() -{ - auto outputParams = std::make_shared(mOutputIds); - - outputParams->sequenceLength = mSeqLengths; - - outputParams->nextDraftTokens = mOutputNextDraftTokens; - - outputParams->numNewTokens = mAcceptedLengths; - - outputParams->nextDraftLengths = mNextDraftLengths; - - outputParams->prevDraftLengths = mPrevDraftLengths; - - outputParams->numNewTokensCumSum = mAcceptedLengthCumSum; - - outputParams->pathsOffsets = mPathsOffsets; - - outputParams->nextDraftPosIds = mNextPosIds; - - outputParams->positionIdsBase = mOutputPositionIdsBase; - - outputParams->randomDataSample = mRandomDataSample; - - outputParams->randomDataValidation = mRandomDataValidation; - - outputParams->packedMasks = mPackedMasks; - - outputParams->packedMasks = mPackedMasks; - - outputParams->unpackedNextDraftTokens = mOutputUnpackedNextDraftTokens; - - outputParams->unpackedNextDraftIndices = mOutputUnpackedNextDraftIndices; - - outputParams->nextDraftProbs = mOutputDraftProbs; - - outputParams->temperatures = mOutputTemperatures; - - outputParams->generationLengths = mOutputGenerationLengths; - - outputParams->generationLengthsHost = mOutputGenerationLengthsHost; - - outputParams->maxGenLengthHost = mMaxGenLengthHost; - - return outputParams; -} - -std::vector boolArrayToBitmask(BufferRange::iterator boolIterator, size_t pathLen) -{ - std::vector bitmask(divUp(pathLen, 32)); - for (size_t bi = 0; bi < pathLen; ++bi) - { - auto slice = bi / 32; - if (boolIterator[bi]) - { - bitmask[slice] |= (1 << (bi % 32)); - } - } - return bitmask; -} - -template -void ExplicitDraftTokensLayerTest::checkLayerResult() -{ - using DataType = typename T::DataType; - auto const batchSlots = BufferRange(*mBatchSlots); - - // Check generated random data - { - auto const randomDataSample = BufferRange(*mRandomDataSample); - auto const randomDataValidation = BufferRange(*mRandomDataValidation); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - // Check that all fields are filled with non zero data - EXPECT_NE(randomDataSample[batchSlot], DataType{0}) << " bi: " << bi; - auto const stride = mSamplingParams.getMaxNumPaths() * mSamplingParams.getMaxDraftPathLen(); - EXPECT_FALSE(std::any_of(randomDataValidation.begin() + batchSlot * stride, - randomDataValidation.begin() + (batchSlot + 1) * stride, - [](DataType val) { return val == DataType{0}; })) - << " bi: " << bi; - } - } - - // Check masks - { - auto const packedMasks = BufferRange(*mPackedMasks); - auto masks = BufferRange(*mMasks); - auto generationLengths = mNetwork.getNextGenerationLengths(); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - for (SizeType32 ti = 0; ti < generationLengths[bi]; ++ti) - { - auto const batchSlot = batchSlots[bi]; - auto const maskIdx = flat_index3( - bi, ti, 0, mNetwork.getMaxNextGenerationLength(), mNetwork.getMaxNextGenerationLength()); - auto const bitmask = boolArrayToBitmask(masks.begin() + maskIdx, mNetwork.getMaxNextGenerationLength()); - for (SizeType32 mi = 0; mi < bitmask.size(); ++mi) - { - auto const packedMaskIdx = flat_index3(batchSlot, ti, mi, mSamplingParams.getMaxDecodingTokens(), - static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))); - EXPECT_EQ(bitmask[mi], packedMasks[packedMaskIdx]) << " bi: " << bi << " ti: " << ti; - } - } - } - } - - // Check accepted tokens - auto const outputIds = BufferRange(*mOutputIds); - auto const refOutputIds = mNetwork.getOutputIds(); - auto const promptIds = mNetwork.getPrompts(); - auto const seqLenghts = BufferRange(*mSeqLengths); - auto const lastDraftTokens = BufferRange(*mLastDraftTokens); - auto const bestPathLengths = BufferRange(*mBestPathLengths); - auto const bestPathIndices = BufferRange(*mBestPathIndices); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - // Updated seq length is prompt length and newly accepted tokens. - EXPECT_EQ(seqLenghts[batchSlot], promptIds[bi].size() + bestPathLengths[bi]) << " bi: " << bi; - // Check that output ids contains accepted tokens. - for (SizeType32 ti = 0; ti < promptIds[bi].size() + bestPathLengths[bi]; ++ti) - { - EXPECT_EQ(outputIds[batchSlot * mSamplingParams.getMaxSeqLen() + ti], refOutputIds[bi][ti]) - << " bi: " << bi << " ti: " << ti; - } - auto outputIter = outputIds.begin() + batchSlot * mSamplingParams.getMaxSeqLen(); - std::vector outputVec(outputIter, outputIter + seqLenghts[batchSlot]); - TLLM_LOG_DEBUG("Output ids at %d request is \"%s\"", bi, mNetwork.detokenize(outputVec).c_str()); - TLLM_LOG_DEBUG("Ref output ids at %d request is \"%s\"", bi, mNetwork.detokenize(refOutputIds[bi]).c_str()); - } - - // Check new draft tokens - { - auto const outputNextDraftTokens = BufferRange(*mOutputNextDraftTokens); - auto const generationLengths = BufferRange(*mSpecDecodingGenerationLengths); - auto const compressedDraftTokens = mNetwork.getNextFlatTokens(); - TLLM_LOG_DEBUG("Next compressed draft tokens are \"%s\"", mNetwork.detokenize(compressedDraftTokens).c_str()); - SizeType32 compressedIdx = 0; - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - auto const generatedLength = generationLengths[bi]; - // Check draft tokens for the next iteration. - for (SizeType32 ti = 0; ti < generatedLength - 1; ++ti) - { - auto const idx = flat_index2(batchSlot, ti, mSamplingParams.getMaxDecodingDraftTokens()); - EXPECT_EQ(outputNextDraftTokens[idx], compressedDraftTokens[compressedIdx + ti + 1]) - << " bi: " << bi << " ti: " << ti; - } - // Check length of the draft tokens. - EXPECT_EQ(BufferRange(*mNextDraftLengths)[batchSlot], generatedLength - 1) << " bi: " << bi; - // Check accepted length. - EXPECT_EQ(BufferRange(*mAcceptedLengths)[batchSlot], bestPathLengths[bi]) << " bi: " << bi; - compressedIdx += generatedLength; - } - } - - // Check position ids - { - auto const outputPositionIdsBase = BufferRange(*mOutputPositionIdsBase); - auto const nextPosIds = BufferRange(*mNextPosIds); - auto const generationLengths = BufferRange(*mSpecDecodingGenerationLengths); - auto const packedPosIds = mNetwork.getNextPackedPosId(); - SizeType32 compressedIdx = 0; - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ(outputPositionIdsBase[batchSlot], seqLenghts[batchSlot]); - - auto const generatedLength = generationLengths[bi]; - // Check pos ids for the next iteration. - for (SizeType32 ti = 0; ti < generatedLength; ++ti) - { - auto const idx = flat_index2(batchSlot, ti, mSamplingParams.getMaxDecodingTokens()); - // Minus -1 to account for context phase correction of pos ids - EXPECT_EQ(nextPosIds[idx], packedPosIds[compressedIdx + ti] - 1) << " bi: " << bi << " ti: " << ti; - } - compressedIdx += generatedLength; - } - } - - // Check unpacked indices and tokens - { - auto const nextDraftTokens = mNetwork.getNextDraftTokens(); - auto const nextDraftIndices = mNetwork.getNextDraftIndices(); - auto const nextDraftTokensRange = BufferRange(*mOutputUnpackedNextDraftTokens); - auto const nextDraftIndicesRange = BufferRange(*mOutputUnpackedNextDraftIndices); - for (SizeType32 bi = 0; bi < nextDraftTokens.size(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - for (SizeType32 pi = 0; pi < nextDraftTokens[bi].size(); ++pi) - { - for (SizeType32 ti = 0; ti < nextDraftTokens[bi][pi].size(); ++ti) - { - auto idx = flat_index3( - batchSlot, pi, ti, mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()); - EXPECT_EQ(nextDraftTokensRange[idx], nextDraftTokens[bi][pi][ti]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti; - EXPECT_EQ(nextDraftIndicesRange[idx], nextDraftIndices[bi][pi][ti]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti; - } - } - } - } - - // Check accumulated cum sum and paths offsets - { - auto const accumulatedCumSum = BufferRange(*mAcceptedLengthCumSum); - auto const pathsOffsets = BufferRange(*mPathsOffsets); - auto const acceptedLengths = BufferRange(*mAcceptedLengths); - auto const bestPathIndices = BufferRange(*mBestPathIndices); - auto const lastDraftIndices = mNetwork.getLastDraftIndices(); - SizeType32 sum = 0; - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ(sum, accumulatedCumSum[bi]) << "bi: " << bi; - auto const acceptedLength = acceptedLengths[batchSlot] - 1; - for (SizeType32 ti = 0; ti < acceptedLength; ++ti) - { - EXPECT_EQ(pathsOffsets[sum + ti], lastDraftIndices[bi][bestPathIndices[bi]][ti + 1] - 1) - << "bi: " << bi << " ti: " << ti; - } - sum += acceptedLength; - } - EXPECT_EQ(sum, accumulatedCumSum[mSamplingParams.getBatchSize()]); - } - - // Check draft probs - { - auto const outDraftProbs = BufferRange(*mOutputDraftProbs); - auto const inDraftProbs = BufferRange(*mNextDraftProbs); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - for (SizeType32 pi = 0; pi < mSamplingParams.getMaxNumPaths(); ++pi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDraftPathLen(); ++ti) - { - for (SizeType32 vi = 0; vi < mSamplingParams.getVocabSize(); ++vi) - { - auto const outProbIdx = flat_index4(batchSlot, pi, ti, vi, mSamplingParams.getMaxNumPaths(), - mSamplingParams.getMaxDraftPathLen(), mSamplingParams.getVocabSize()); - auto const inProbIdx = flat_index4(bi, pi, ti, vi, mSamplingParams.getMaxNumPaths(), - mSamplingParams.getMaxDraftPathLen(), mSamplingParams.getVocabSize()); - EXPECT_EQ(outDraftProbs[outProbIdx], inDraftProbs[inProbIdx]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti << " vi: " << vi; - } - } - } - } - } - - // Check temperature - { - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ( - BufferRange(*mOutputTemperatures)[batchSlot], static_cast(1.f / mTemperatures[bi])) - << " bi: " << bi; - } - } -} - -template -void ExplicitDraftTokensLayerTest::packData() -{ - using DataType = typename T::DataType; - tksd::PackExplicitDraftTokensParams params; - params.batchSlots = bufferCast(*mBatchSlots); - params.cumSumGenerationLengths = bufferCast(*mCumSumGenerationLengths); - params.maxGenerationLength = bufferCast(*mMaxGenerationLength); - - params.outputPositionIdsBase = bufferCast(*mPackedPositionIdsBase); - params.inputPositionIdsBase = bufferCast(*mOutputPositionIdsBase); - - params.outputGenerationLengths = bufferCast(*mPackedGenerationLengths); - params.inputGenerationLengths = bufferCast(*mSpecDecodingGenerationLengths); - - params.outputRandomDataSample = bufferCast(*mPackedRandomDataSample); - params.inputRandomDataSample = bufferCast(*mRandomDataSample); - - params.outputRandomDataValidation = bufferCast(*mPackedRandomDataVerification); - params.inputRandomDataValidation = bufferCast(*mRandomDataValidation); - - params.outputNextDraftTokens = bufferCast(*mPackedNextDraftTokens); - params.inputNextDraftTokens = bufferCast(*mOutputUnpackedNextDraftTokens); - - params.outputNextDraftIndices = bufferCast(*mPackedNextDraftIndices); - params.inputNextDraftIndices = bufferCast(*mOutputUnpackedNextDraftIndices); - - params.outputPackedMask = bufferCast(*mPackedPackedMasks); - params.inputPackedMask = bufferCast(*mPackedMasks); - - params.inputPositionIds = bufferCast(*mNextPosIds); - params.outputPositionOffsets = bufferCast(*mPackedPositionOffsets); - params.outputPositionIds = bufferCast(*mPackedPackedPosIds); - - params.outputDraftProbs = bufferCast(*mPackedDraftProbs); - params.inputDraftProbs = bufferCast(*mOutputDraftProbs); - - params.outputTemperatures = bufferCast(*mPackedTemperatures); - params.inputTemperatures = bufferCast(*mOutputTemperatures); - - params.batchSize = mSamplingParams.getBatchSize(); - params.numPaths = mSamplingParams.getMaxNumPaths(); - params.maxPathLength = mSamplingParams.getMaxPathLen(); - params.vocabSize = mSamplingParams.getVocabSize(); - params.numGenerationRequests = mSamplingParams.getBatchSize(); - params.numContextTokens = 0; - - params.checkParams(); - - tksd::invokePackGenerationLengths(params, mStream->get()); - - // Compute inclusive sum - auto reduceTempStorageBytes = tksd::invokeScanGenerationLengths( - nullptr, 0, nullptr, nullptr, mSamplingParams.getBatchSize(), mStream->get()); - auto reduceMaxTempStorage = mBufferManager->gpu(reduceTempStorageBytes); - tksd::invokeScanGenerationLengths(bufferCast(*reduceMaxTempStorage), reduceTempStorageBytes, - bufferCast(*mSpecDecodingGenerationLengths), bufferCast(*mCumSumGenerationLengths), - mSamplingParams.getBatchSize(), mStream->get()); - - // Pack tensors from batch slot position to continuous array - tksd::invokePackExplicitDraftTokens(params, mStream->get()); - - // Copy draft probs - tksd::invokeCopyProbs(params, mStream->get()); -} - -template -void ExplicitDraftTokensLayerTest::checkPackResult() -{ - using DataType = typename T::DataType; - auto const batchSlots = BufferRange(*mBatchSlots); - auto const maxGenLength = mNetwork.getMaxNextGenerationLength(); - auto const numPackedMasks = static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32)); - for (SizeType32 bi = 0; bi < mSamplingParams.getBatchSize(); ++bi) - { - auto const batchSlot = batchSlots[bi]; - EXPECT_EQ(BufferRange(*mPackedPositionIdsBase)[bi], - BufferRange(*mOutputPositionIdsBase)[batchSlot]) - << "bi: " << bi; - EXPECT_EQ(BufferRange(*mPackedGenerationLengths)[bi], - BufferRange(*mSpecDecodingGenerationLengths)[batchSlot]) - << "bi: " << bi; - EXPECT_EQ( - BufferRange(*mPackedRandomDataSample)[bi], BufferRange(*mRandomDataSample)[batchSlot]) - << "bi: " << bi; - EXPECT_EQ( - BufferRange(*mPackedTemperatures)[bi], BufferRange(*mOutputTemperatures)[batchSlot]) - << "bi: " << bi; - - for (SizeType32 pi = 0; pi < mSamplingParams.getMaxNumPaths(); ++pi) - { - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxDraftPathLen(); ++ti) - { - EXPECT_EQ(bufferCast(*ITensor::at(mPackedRandomDataVerification, {bi, pi, ti}))[0], - bufferCast(*ITensor::at(mRandomDataValidation, {batchSlot, pi, ti}))[0]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti; - for (SizeType32 vi = 0; vi < mSamplingParams.getVocabSize(); ++vi) - { - EXPECT_EQ(bufferCast(*ITensor::at(mPackedDraftProbs, {bi, pi, ti, vi}))[0], - bufferCast(*ITensor::at(mOutputDraftProbs, {batchSlot, pi, ti, vi}))[0]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti << " vi: " << vi; - } - } - for (SizeType32 ti = 0; ti < mSamplingParams.getMaxPathLen(); ++ti) - { - EXPECT_EQ(bufferCast(*ITensor::at(mPackedNextDraftTokens, {bi, pi, ti}))[0], - bufferCast(*ITensor::at(mOutputUnpackedNextDraftTokens, {batchSlot, pi, ti}))[0]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti; - EXPECT_EQ(bufferCast(*ITensor::at(mPackedNextDraftIndices, {bi, pi, ti}))[0], - bufferCast(*ITensor::at(mOutputUnpackedNextDraftIndices, {batchSlot, pi, ti}))[0]) - << "bi: " << bi << " pi: " << pi << " ti: " << ti; - } - } - auto const basePosId = BufferRange(*mPackedPositionIdsBase)[bi]; - for (SizeType32 ti = 0; ti < maxGenLength; ++ti) - { - auto const outPosOffsetIdx = flat_index2(bi, ti, maxGenLength); - auto const inPosOffsetIdx = flat_index2(batchSlot, ti, mSamplingParams.getMaxDecodingTokens()); - EXPECT_EQ(BufferRange(*mPackedPositionOffsets)[outPosOffsetIdx], - BufferRange(*mNextPosIds)[inPosOffsetIdx] - basePosId + 1) - << "bi: " << bi << " ti: " << ti; - } - auto const outputMaskStartId = (bi == 0) ? 0 : BufferRange(*mCumSumGenerationLengths)[bi - 1]; - auto const numTokens = (bi == 0) ? BufferRange(*mCumSumGenerationLengths)[0] - : BufferRange(*mCumSumGenerationLengths)[bi] - - BufferRange(*mCumSumGenerationLengths)[bi - 1]; - for (SizeType32 mi = 0; mi < numTokens * numPackedMasks; ++mi) - { - auto const outMaskIdx = outputMaskStartId * numPackedMasks + mi; - auto const inMaskIdx = flat_index2(batchSlot, mi, mSamplingParams.getMaxDecodingTokens() * numPackedMasks); - EXPECT_EQ( - BufferRange(*mPackedPackedMasks)[outMaskIdx], BufferRange(*mPackedMasks)[inMaskIdx]) - << "bi: " << bi << " mi: " << mi; - } - } -} - -template -void ExplicitDraftTokensLayerTest::runTest(std::vector const& prompts, - std::vector const& predictions, DraftLettersVec const& nextDraftLetters, - DraftLettersVec const& lastDraftLetters, SamplingParams& params) -{ - mSamplingParams = params; - - mNetwork.forward(params, prompts, predictions, nextDraftLetters, lastDraftLetters); - - allocateBuffers(); - - setup(); - - auto inputTensors = createInputTensors(); - auto outputTensors = createOutputTensors(); - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mExplicitDraftTokensLayer->forwardAsync(outputTensors, inputTensors, mDecodingWorkspace); - - mStream->synchronize(); - - checkLayerResult(); - - packData(); - - mStream->synchronize(); - - checkPackResult(); -} - -template class ExplicitDraftTokensLayerTest>; -template class ExplicitDraftTokensLayerTest>; -#ifdef ENABLE_BF16 -template class ExplicitDraftTokensLayerTest>; -#endif // ENABLE_BF16 - -TYPED_TEST_SUITE(ExplicitDraftTokensLayerTest, TestTypes); - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestBS1) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h"}; - std::vector predictions = {"how things"}; - DraftLettersVec lastDraftLetters = {{"how do ", "how are", "however", "hello w"}}; - DraftLettersVec nextDraftLetters = {{"things ", "that is", "to crea", "touchab"}}; - - params.setBatchSize(1); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestBS1OnePaths) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h"}; - std::vector predictions = {"how things"}; - DraftLettersVec lastDraftLetters = {{"how do "}}; - DraftLettersVec nextDraftLetters = {{"things "}}; - - params.setBatchSize(1); - params.setMaxNumPaths(1); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestSecondPathAcceptedBS1) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h"}; - std::vector predictions = {"how things"}; - DraftLettersVec lastDraftLetters = {{"howdy f", "how are", "however", "hello w"}}; - DraftLettersVec nextDraftLetters = {{"things ", "that is", "to crea", "touchab"}}; - - params.setBatchSize(1); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestNoDraftAcceptedBS1) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h"}; - std::vector predictions = {"how things"}; - DraftLettersVec lastDraftLetters = {{"handove", "human f", "heavy l", "hello h"}}; - DraftLettersVec nextDraftLetters = {{"oatmeal", "ocean b", "occupat", "oblivio"}}; - - params.setBatchSize(1); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestBS2SameSequence) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h", "Hi mate, h"}; - std::vector predictions = {"how things", "how things"}; - DraftLettersVec lastDraftLetters - = {{"how do ", "how are", "however", "hello w"}, {"how do ", "how are", "however", "hello w"}}; - DraftLettersVec nextDraftLetters - = {{"things ", "that is", "to crea", "touchab"}, {"things ", "that is", "to crea", "touchab"}}; - - params.setBatchSize(2); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestBS2Long) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h", "London is t"}; - std::vector predictions = {"how things are going", "the capital of Great Britain"}; - DraftLettersVec lastDraftLetters = {{"how do you ", "how are you", "however you", "hello world"}, - {"the bar and", "the best ci", "the capital", "thoughest p"}}; - DraftLettersVec nextDraftLetters = {{"things are ", "that is sad", "to create a", "touchable y"}, - {" of Great B", " and the ma", " of country", " also known"}}; - - params.setBatchSize(2); - // ceil(4 * 10 / 32) = 2 masks per request - params.setMaxNumPaths(4); - params.setMaxDraftPathLen(10); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestBS2DifferentSequences) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h", "London is t"}; - std::vector predictions = {"how things", "the cap"}; - DraftLettersVec lastDraftLetters - = {{"how do ", "how are", "however", "hello w"}, {"the bar", "the bes", "the cap", "thoughe"}}; - DraftLettersVec nextDraftLetters - = {{"things ", "that is", "to crea", "touchab"}, {"itan of", "iteract", "ital of", "importa"}}; - - params.setBatchSize(2); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -TYPED_TEST(ExplicitDraftTokensLayerTest, SimpleTestB4DifferentSequences) -{ - SamplingParams params; - - std::vector prompt = {"Hi mate, h", "London is t", "Short", "Very long prompt but should not m"}; - std::vector predictions = {"how things", "the cap", "twave o", "matter "}; - DraftLettersVec lastDraftLetters - = {{"how do ", "how are", "however", "hello w"}, {"the bar", "the bes", "the cap", "thoughe"}, - {"t promp", "ts on Y", "ter out", "twave o"}, {"matter ", "mean an", "make th", "modify "}}; - DraftLettersVec nextDraftLetters - = {{"things ", "that is", "to crea", "touchab"}, {"itan of", "iteract", "ital of", "importa"}, - {" chips ", " oil an", " semico", " exampl"}, {"at all ", "anythin", "above a", "albeit "}}; - - params.setBatchSize(4); - - this->runTest(prompt, predictions, nextDraftLetters, lastDraftLetters, params); -} - -template -class FillRandDataTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - static auto constexpr mDataType{TRTDataType::value}; - - FillRandDataTest() {} - - void SetUp() override - { - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); - } - - void TearDown() override {} - - void runTest(SizeType32 batchSize, SizeType32 numPaths, SizeType32 draftLength, bool skipVerification, - uint64_t randomSeed, bool batchInit) - { - SizeType32* batchSlotsPtr{nullptr}; - - auto curandState = mBufferManager->gpu(ITensor::makeShape({batchSize, 48}), tensorrt_llm::DataType::kUINT8); - auto* curandStatePtr = reinterpret_cast(bufferCast(*curandState)); - - if (batchInit) - { - auto randomSeeds = mBufferManager->gpu(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT64); - trk::invokeFill(*randomSeeds, static_cast(randomSeed), *mStream); - auto* randomSeedsPtr = bufferCast(*randomSeeds); - tk::invokeCurandBatchInitialize(curandStatePtr, batchSlotsPtr, batchSize, randomSeedsPtr, mStream->get()); - } - else - { - tk::invokeCurandInitialize(curandStatePtr, batchSlotsPtr, batchSize, randomSeed, mStream->get()); - } - mStream->synchronize(); - - tksd::FillRandDataExplicitDraftTokensParams params; - params.batchSize = batchSize; - params.numPaths = numPaths; - params.draftLength = draftLength; - params.skipVerification = skipVerification; - - auto randDataSample = mBufferManager->gpu(ITensor::makeShape({batchSize}), mDataType); - auto randDataValidation - = mBufferManager->gpu(ITensor::makeShape({batchSize, numPaths, draftLength}), mDataType); - - params.randDataSample = bufferCast(*randDataSample); - params.randDataVerification = bufferCast(*randDataValidation); - params.curandState = curandStatePtr; - params.batchSlots = batchSlotsPtr; - - tksd::invokeFillRandData(params, mStream->get()); - mStream->synchronize(); - - auto randDataSampleHost = mBufferManager->copyFrom(*randDataSample, MemoryType::kCPU); - auto randDataSampleHostPtr = bufferCast(*randDataSampleHost); - EXPECT_GE(randDataSampleHostPtr[0], T(0)); - EXPECT_LE(randDataSampleHostPtr[0], T(1)); - - auto randDataValidationHost = mBufferManager->copyFrom(*randDataValidation, MemoryType::kCPU); - auto randDataValidationHostRange = BufferRange(*randDataValidationHost); - for (auto i = 0; i < randDataValidationHostRange.size(); ++i) - { - EXPECT_GE(randDataValidationHostRange[i], T(0)) << "index " << i; - EXPECT_LE(randDataValidationHostRange[i], T(1)) << "index " << i; - } - } - -private: - std::shared_ptr mStream; - std::shared_ptr mBufferManager; -}; - -#ifdef ENABLE_BF16 -using FloatHalfBfloatTypes = testing::Types; -TYPED_TEST_SUITE(FillRandDataTest, FloatHalfBfloatTypes); -#else -TYPED_TEST_SUITE(FillRandDataTest, FloatAndHalfTypes); -#endif - -TYPED_TEST(FillRandDataTest, SimpleTest) -{ - SizeType32 constexpr batchSize{2}; - SizeType32 constexpr numPaths{3}; - SizeType32 constexpr draftLength{4}; - bool constexpr skipVerification{false}; - - uint64_t randomSeed{0}; - - this->runTest(batchSize, numPaths, draftLength, skipVerification, randomSeed, false); -} - -TYPED_TEST(FillRandDataTest, BatchInit) -{ - SizeType32 constexpr batchSize{3}; - SizeType32 constexpr numPaths{2}; - SizeType32 constexpr draftLength{5}; - bool constexpr skipVerification{false}; - - uint64_t randomSeed{42}; - - this->runTest(batchSize, numPaths, draftLength, skipVerification, randomSeed, true); -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.h b/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.h deleted file mode 100644 index a956bca97b03..000000000000 --- a/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.h +++ /dev/null @@ -1,351 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "tensorrt_llm/layers/explicitDraftTokensLayer.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -#include - -#include - -namespace tensorrt_llm::tests::layers -{ - -class SamplingParams -{ -public: - SamplingParams() {} - - inline void setBatchSize(runtime::SizeType32 batchSize) - { - mBatchSize = batchSize; - } - - inline void setMaxNumPaths(runtime::SizeType32 maxNumPaths) - { - mMaxNumPaths = maxNumPaths; - } - - inline void setMaxDraftPathLen(runtime::SizeType32 maxDraftPathLen) - { - mMaxDraftPathLen = maxDraftPathLen; - } - - [[nodiscard]] inline runtime::SizeType32 getBatchSize() const - { - return mBatchSize; - } - - [[nodiscard]] inline runtime::SizeType32 getVocabSize() const - { - return mVocabSize; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxBatchSize() const - { - return 2 * getBatchSize(); - } - - [[nodiscard]] inline runtime::SizeType32 getMaxDraftPathLen() const - { - return mMaxDraftPathLen; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxPathLen() const - { - return getMaxDraftPathLen() + 1; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxNumPaths() const - { - return mMaxNumPaths; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxDecodingDraftTokens() const - { - return getMaxDraftPathLen() * getMaxNumPaths(); - } - - [[nodiscard]] inline runtime::SizeType32 getMaxDecodingTokens() const - { - return getMaxDecodingDraftTokens() + 1; - } - - [[nodiscard]] inline runtime::SizeType32 getMaxSeqLen() const - { - return getMaxDecodingTokens() * 2; - } - - [[nodiscard]] inline runtime::TokenIdType getPadId() const - { - return mPadId; - } - -private: - runtime::SizeType32 mBatchSize{6}; - runtime::SizeType32 mMaxDraftPathLen{6}; - runtime::SizeType32 mMaxNumPaths{4}; - runtime::TokenIdType mPadId{-1}; - runtime::SizeType32 mVocabSize{256}; -}; - -using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; -using BufferPtr = tensorrt_llm::runtime::IBuffer::SharedPtr; -using SizeType32 = tensorrt_llm::runtime::SizeType32; -using TokenIdType = tensorrt_llm::runtime::TokenIdType; - -using TokensVec = std::vector; -using DraftLettersVec = std::vector>; -using DraftTokensVec = std::vector>; -using DraftTokensIndices = std::vector>>; - -class ExplicitDraftTokensDummyNetwork -{ -public: - void forward(SamplingParams const& params, std::vector const& prompts, - std::vector const& predictionLetters, DraftLettersVec const& nextDraftLetters, - DraftLettersVec const& lastDraftLetters); - - TokensVec tokenize(std::string const& letters) const; - - std::string detokenize(TokensVec const& tokens) const; - - DraftTokensVec draftLettersToTokens(DraftLettersVec const& draftLetters) const; - - SizeType32 longestCommonPrefixLength(TokensVec const& a, TokensVec const& b) const; - - SizeType32 computeCompressedVectorAndIndices(TokensVec& compressedVector, std::vector& packedPosIds, - DraftTokensIndices& indices, std::vector const& vectors, SizeType32 basePosId); - - void compressTokens(TokensVec& compressedVector, std::vector& packedPosIds, DraftTokensIndices& indices, - std::vector& generationLengths, DraftTokensVec const& draftTokens, - std::vector const& basePosIds); - - void acceptTokens(std::vector const& predictionTokens, DraftTokensVec const& lastDraftTokens, - DraftTokensVec const& nextDraftTokens); - - void createNextMasks(DraftTokensIndices const& indices, DraftTokensVec const& draftTokens, SizeType32 maxGenLength); - - void setSamplingParams(SamplingParams const& params) - { - mSamplingParams = params; - } - - std::vector getPrompts() const - { - return mPrompts; - } - - std::vector getOutputIds() const - { - return mOutputIds; - } - - TokensVec getNextFlatTokens() const - { - return mNextCompressedVector; - } - - DraftTokensVec getNextDraftTokens() const - { - return mNextDraftTokens; - } - - DraftTokensIndices getNextDraftIndices() const - { - return mNextDraftTokenIndices; - } - - DraftTokensIndices getLastDraftIndices() const - { - return mLastDraftTokenIndices; - } - - DraftTokensVec getLastDraftTokens() const - { - return mLastDraftTokens; - } - - std::vector getBestPathLengths() const - { - return mBestPathLengths; - } - - std::vector getBestPathIndices() const - { - return mBestPathIndices; - } - - std::vector getNextPackedPosId() const - { - return mNextPackedPosIds; - } - - std::vector getNextGenerationLengths() const - { - return mNextGenerationLengths; - } - - SizeType32 getMaxNextGenerationLength() const - { - return mMaxNextGenLength; - } - - std::vector>> getNextMasks() const - { - return mMasks; - } - -private: - SamplingParams mSamplingParams; - - std::vector mPrompts; - std::vector mOutputIds; - - DraftTokensVec mNextDraftTokens; - DraftTokensVec mLastDraftTokens; - - TokensVec mNextCompressedVector; - std::vector mNextPackedPosIds; - DraftTokensIndices mNextDraftTokenIndices; - - TokensVec mLastCompressedVector; - std::vector mLastPackedPosIds; - DraftTokensIndices mLastDraftTokenIndices; - - std::vector mBestPathLengths; - std::vector mBestPathIndices; - - std::vector mNextGenerationLengths; - std::vector mLastGenerationLengths; - SizeType32 mMaxNextGenLength; - - std::vector>> mMasks; -}; - -template -class ExplicitDraftTokensLayerTest : public testing::Test -{ -private: - void SetUp() override; - -private: - SamplingParams mSamplingParams; - - // Outputs - TensorPtr mSeqLengths; - TensorPtr mAcceptedLengths; - TensorPtr mOutputIds; - TensorPtr mOutputNextDraftTokens; - TensorPtr mOutputPositionIdsBase; - TensorPtr mRandomDataSample; - TensorPtr mRandomDataValidation; - TensorPtr mAcceptedLengthCumSum; - TensorPtr mPackedMasks; - TensorPtr mPathsOffsets; - TensorPtr mNextPosIds; - TensorPtr mNextDraftLengths; - TensorPtr mPrevDraftLengths; - TensorPtr mOutputUnpackedNextDraftTokens; - TensorPtr mOutputUnpackedNextDraftIndices; - TensorPtr mOutputDraftProbs; - TensorPtr mOutputTemperatures; - TensorPtr mOutputGenerationLengths; - TensorPtr mOutputGenerationLengthsHost; - TensorPtr mMaxGenLengthHost; - - // inputs - TensorPtr mBatchSlots; - TensorPtr mMasks; - TensorPtr mInputNextDraftTokens; - TensorPtr mNextDraftIndices; - TensorPtr mLastDraftTokens; - TensorPtr mLastDraftIndices; - TensorPtr mNextDraftProbs; - TensorPtr mPackedPosIds; - TensorPtr mBestPathLengths; - TensorPtr mBestPathIndices; - TensorPtr mSpecDecodingGenerationLengths; - TensorPtr mTokensPerStep; - TensorPtr mNextFlatTokens; - TensorPtr mInputPositionIdsBase; - TensorPtr mEndIds; - TensorPtr mMaxGenLengthDevice; - - // Packed inputs - TensorPtr mMaxGenerationLength; - TensorPtr mCumSumGenerationLengths; - - // Packed outputs - TensorPtr mPackedPositionIdsBase; - TensorPtr mPackedGenerationLengths; - TensorPtr mPackedRandomDataSample; - TensorPtr mPackedRandomDataVerification; - TensorPtr mPackedNextDraftTokens; - TensorPtr mPackedNextDraftIndices; - TensorPtr mPackedPackedMasks; - TensorPtr mPackedPositionOffsets; - TensorPtr mPackedPackedPosIds; - TensorPtr mPackedDraftProbs; - TensorPtr mPackedTemperatures; - - // Setup params - std::vector mRandomSeeds; - std::vector mTemperatures; - - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - std::shared_ptr> mExplicitDraftTokensLayer; - std::shared_ptr mDecodingWorkspace; - - ExplicitDraftTokensDummyNetwork mNetwork; - -private: - void allocateBuffers(); - - void setup(); - - std::shared_ptr createInputTensors(); - - std::shared_ptr createOutputTensors(); - - void checkLayerResult(); - - void packData(); - - void checkPackResult(); - -public: - void runTest(std::vector const& prompts, std::vector const& predictions, - DraftLettersVec const& nextDraftLetters, DraftLettersVec const& lastDraftLetters, SamplingParams& params); -}; - -template -struct TypePair -{ - using LayerType = T; - using DataType = U; -}; - -#ifdef ENABLE_BF16 -using TestTypes = testing::Types, TypePair, TypePair>; -#else -using TestTypes = testing::Types, TypePair>; -#endif // ENABLE_BF16 - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp b/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp deleted file mode 100644 index cdc913d3b3aa..000000000000 --- a/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp +++ /dev/null @@ -1,1408 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tests/unit_tests/layers/baseSamplingLayerTest.h" - -#include - -namespace -{ - -namespace tle = tensorrt_llm::executor; -namespace trk = tensorrt_llm::runtime::kernels; - -using namespace tensorrt_llm::tests::layers::sampling; -using namespace tensorrt_llm::layers; -using namespace tensorrt_llm::runtime; - -template -class ExternalDraftTokensLayerTest : public BaseSamplingLayerTest -{ -protected: - int32_t const mMaxDraftLen = this->mMaxTokensPerEngineStep - 1; - - TensorPtr mDraftLogits; - TensorPtr mDraftProbs; - TensorPtr mTargetProbs; - TensorPtr mNumDraftTokens; - TensorPtr mNumDraftTokensHost; - TensorPtr mDraftTokenIds; - TensorPtr mUseDraftLogits; - TensorPtr mUseDraftLogitsHost; - float mConstantThreshold = 1.0f; - bool mUseRandomAcceptanceThreshold = true; - - std::vector* mTestDraftLogitsInit; - std::vector mTestDraftLogitsAccept = { - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, // step 0 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // step 1 - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // step 2 - }; - std::vector mTestDraftLogitsReject = { - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, // step 0 - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // step 1 - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // step 2 - }; - std::vector> mTestDraftTokenIdsInit; - - void SetUp() override - { - this->mStream = std::make_shared(); - this->mBufferManager = std::make_shared(this->mStream); - } - - void initLayer(TestSamplingParams const& params) override - { - auto decodingMode = tle::DecodingMode::ExternalDraftTokens(); - - auto const decodingDomain - = tensorrt_llm::layers::DecoderDomain(this->maxBatchSize(), 1, this->mVocabSize, this->mVocabSizePadded); - this->mSamplingLayer = std::make_shared>( - decodingMode, decodingDomain, this->mBufferManager, true, params.isAirTopPExternalDraftTokensLayer); - - auto const dataType = TRTDataType::value; - - mDraftLogits = this->mBufferManager->gpu( - ITensor::makeShape({this->maxBatchSize(), mMaxDraftLen, this->mVocabSize}), dataType); - mDraftProbs = this->mBufferManager->gpu( - ITensor::makeShape({this->maxBatchSize(), mMaxDraftLen, this->mBeamWidth, this->mVocabSize}), dataType); - mTargetProbs = this->mBufferManager->gpu( - ITensor::makeShape( - {this->maxBatchSize(), this->mMaxTokensPerEngineStep, this->mBeamWidth, this->mVocabSize}), - dataType); - - mDraftTokenIds = this->mBufferManager->gpu( - ITensor::makeShape({this->maxBatchSize(), mMaxDraftLen}), tensorrt_llm::DataType::kINT32); - mUseDraftLogits - = this->mBufferManager->gpu(ITensor::makeShape({this->maxBatchSize()}), TRTDataType::value); - mUseDraftLogitsHost - = this->mBufferManager->cpu(ITensor::makeShape({this->maxBatchSize()}), TRTDataType::value); - - mNumDraftTokens - = this->mBufferManager->gpu(ITensor::makeShape({this->maxBatchSize()}), TRTDataType::value); - mNumDraftTokensHost - = this->mBufferManager->cpu(ITensor::makeShape({this->maxBatchSize()}), TRTDataType::value); - - batchCopyDraftTokenIds(); - if (params.useDraftLogits) - { - batchCopyDraftLogits(); - } - - batchUseDraftLogits(params.useDraftLogits); - } - - std::shared_ptr createInputTensors(int32_t step) override - { - constexpr int32_t ite = 0; - auto decodeInputTensors = std::make_shared( - this->mEndIdsDevice, this->mBatchSlots, step, ite, this->mBatchSize); - - decodeInputTensors->logits = this->mDecodingWorkspace->getDeviceRuntimeLogits(); - - decodeInputTensors->inputLengths = this->mContextLengthDevice; - - decodeInputTensors->finished = this->mFinishedDevice; - - decodeInputTensors->probsComputed = this->mComputeProbs; - - decodeInputTensors->curandStates - = reinterpret_cast(bufferCast(*this->mCurandStatesDevice)); - - decodeInputTensors->draftLogits = mDraftLogits; - decodeInputTensors->draftProbs = mDraftProbs; - decodeInputTensors->targetProbs = mTargetProbs; - decodeInputTensors->numDraftTokens = mNumDraftTokens; - decodeInputTensors->numDraftTokensHost = mNumDraftTokensHost; - decodeInputTensors->draftTokenIds = mDraftTokenIds; - decodeInputTensors->constantThreshold = mConstantThreshold; - decodeInputTensors->useRandomAcceptanceThreshold = mUseRandomAcceptanceThreshold; - decodeInputTensors->step = step; - decodeInputTensors->useDraftLogits = mUseDraftLogits; - decodeInputTensors->useDraftLogitsHost = mUseDraftLogitsHost; - - return decodeInputTensors; - } - - void batchCopyDraftLogits(); - void batchCopyDraftTokenIds(); - void batchUseDraftLogits(bool useDraftLogits); -}; - -template -void ExternalDraftTokensLayerTest::batchCopyDraftLogits() -{ - auto const draftLogitsHost = ITensor::wrap( - mTestDraftLogitsInit->data(), TRTDataType::value, ITensor::makeShape({mMaxDraftLen, this->mVocabSize})); - TLLM_CHECK(mTestDraftLogitsInit->size() == draftLogitsHost->getSize()); - - for (int32_t bi = 0; bi < this->mBatchSize; ++bi) - { - auto draftLogitsDeviceView - = ITensor::slice(mDraftLogits, bi * ExternalDraftTokensLayerTest::kDoubleBatchIdx, 1); - this->mBufferManager->copy(*draftLogitsHost, *draftLogitsDeviceView); - } -} - -template -void ExternalDraftTokensLayerTest::batchCopyDraftTokenIds() -{ - auto numDraftTokensHostRange = BufferRange(*mNumDraftTokensHost); - - for (int32_t bi = 0; bi < this->mBatchSize; ++bi) - { - auto batchSlot = bi * ExternalDraftTokensLayerTest::kDoubleBatchIdx; - - auto const& draftTokenIdsHost = mTestDraftTokenIdsInit.at(bi); - numDraftTokensHostRange[batchSlot] = draftTokenIdsHost.size(); - - auto draftTokenIdsDeviceView = ITensor::at(mDraftTokenIds, {batchSlot}); - TLLM_CHECK(draftTokenIdsDeviceView->getSize() == mMaxDraftLen); - - draftTokenIdsDeviceView->resize(draftTokenIdsHost.size()); - TLLM_CHECK(draftTokenIdsDeviceView->getSize() == draftTokenIdsHost.size()); - - this->mBufferManager->copy(draftTokenIdsHost.data(), *draftTokenIdsDeviceView); - } - - this->mBufferManager->copy(*this->mNumDraftTokensHost, *this->mNumDraftTokens); -} - -template -void ExternalDraftTokensLayerTest::batchUseDraftLogits(bool useDraftLogits) -{ - auto useDraftLogitsHost = BufferRange(*this->mUseDraftLogitsHost); - std::fill(useDraftLogitsHost.begin(), useDraftLogitsHost.end(), useDraftLogits); - trk::invokeFill(*this->mUseDraftLogits, useDraftLogits, *this->mStream); -} - -TYPED_TEST_SUITE(ExternalDraftTokensLayerTest, FloatAndHalfTypes); - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopK) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - // step 0, only token 4 and 5 (topK==2) get accepted - // step 1, only token 0 and 1 gets accepted - // step 2, only token 2 and 3 gets accepted - // step 3, bonus step, token 0 and 1 can be sampled - this->mTestDraftTokenIdsInit = { - {4, 1, 2}, // - {5, 0, 3}, // - {4, 1, 2}, // - {5, 0, 3}, // - {4, 1, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {5}, {4}, {5}, {4}, {4, 5}, // step 0 - {1}, {0}, {1}, {0}, {1}, {0}, // step 1 - {2}, {3}, {2}, {3}, {2}, {0}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopKReject) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 2}, // accept, accept, accept, sampled - {4, 3, 4}, // accept, reject, 0, 0 - {4, 3, 4}, // accept, reject, 0, 0 - {2, 3, 4}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4, 5, 6, 7}, {4, 5}, // step 0 - {0}, {0}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0, 1}, {0, 1}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopK1TopP0) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopK1TopP0Reject) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 3}, // accept, accept, reject, 0 - {4, 1, 2}, // accept, reject, 0, 0 - {4, 1, 2}, // accept, reject, 0, 0 - {5, 0, 2}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsBatchTopK) -{ - std::vector topKs = {1, 1, 2, 2, 4, 4}; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topKs}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4, 5, 6, 7}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0, 1}, {0, 1}, {0, 1, 2, 3}, {0, 1, 2, 3}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsBatchTopKReject) -{ - std::vector topKs = {1, 1, 2, 2, 4, 4}; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topKs}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {}, // no draft tokens, token will be sampled - {5, 0, 2}, // reject, 0, 0, 0 - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 4}, // accept, accept, reject, 0 - {4, 0, 2}, // accept, accept, accept, sampled - {4, 5, 2}, // accept, reject, 0, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4, 5, 6, 7}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0, 1, 2, 3}, // step 1 - {0}, {0}, {2}, {2, 3}, {2}, {0}, // step 2 - {0}, {0}, {0, 1}, {0}, {0, 1, 2, 3}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopP) -{ - // Skip topK decode - float topP = 0.3; - TestSamplingParams params; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopPReject) -{ - // Skip topK decode - float topP = 0.3; - TestSamplingParams params; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 2}, // accept, accept, reject, 0 - {4, 1, 3}, // accept, reject, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopKTopP) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsTopKTopPReject) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 3}, // accept, accept, reject, 0 - {4, 3, 2}, // accept, reject, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsBatchTopKBatchTopP) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {6, 2, 4}, // - {4, 0, 2}, // - {4, 0, 2}, // - {5, 1, 3}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {6}, {4}, {4}, {5}, {4}, {4}, // step 0 - {2}, {0}, {0}, {1}, {0}, {0}, // step 1 - {4}, {2}, {2}, {3}, {2}, {0}, // step 2 - {0, 1, 2}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsBatchTopKBatchTopPReject) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {6, 3, 4}, // accept, reject, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - {4, 2, 2}, // accept, reject, 0, 0 - {7, 1, 3}, // reject, 0, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {6}, {4}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {2}, {0}, {0}, {2}, {0}, // step 2 - {0, 1, 2}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsBatchTopK0BatchTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {7, 3, 5}, // - {5, 1, 3}, // - {5, 1, 3}, // - {5, 1, 3}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {7}, {5}, {5}, {5}, {4}, {4}, // step 0 - {3}, {1}, {1}, {1}, {0}, {0}, // step 1 - {5}, {3}, {3}, {3}, {2}, {0}, // step 2 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByLogitsBatchTopK0BatchTopPReject) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {7, 4, 5}, // accept, reject, 0, 0 - {5, 5, 3}, // accept, reject, 0, 0 - {6, 1, 3}, // reject, 0, 0, 0 - {5, 1, 3}, // accept, accept, accept, sampled - {4, 2, 2}, // accept, reject, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {7}, {5}, {4, 5}, {5}, {4}, {4}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0}, {1}, {0}, {0}, // step 1 - {0}, {0}, {0}, {3}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByTokenIdsBatchTopKBatchTopP) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {0, 2, 3, 4}, {2}, {2}, {0, 2, 3}, {2}, {0}, // step 2 - {0, 1, 2}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByTokenIdsBatchTopKBatchTopPReject) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {4, 3, 2}, // accept, reject, 0, 0 - {5, 0, 2}, // reject, 0, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - {6, 0, 2}, // reject, 0, 0, 0 - {4, 1, 2}, // accept, reject, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {0}, {2}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByTokenIdsBatchTopK0BatchTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {7, 3, 5}, // - {5, 1, 3}, // - {5, 1, 3}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0}, // step 1 - {0, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 2, 3}, {0, 2, 3}, {2}, {0}, // step 2 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AcceptByTokenIdsBatchTopK0BatchTopPReject) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {7, 4, 5}, // accept, reject, 0, 0 - {5, 1, 6}, // accept/reject, accept/reject, reject, 0 - {6, 1, 3}, // reject, 0, 0, 0 - {4, 0, 2}, // accept/reject, accept/reject, accept, sampled - {4, 1, 2}, // accept, reject, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0}, {0, 1}, {0}, {0}, // step 1 - {0}, {0, 2, 3, 4, 5}, {0}, {0, 2, 3}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopK) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - // step 0, only token 4 and 5 (topK==2) get accepted - // step 1, only token 0 and 1 gets accepted - // step 2, only token 2 and 3 gets accepted - // step 3, bonus step, token 0 and 1 can be sampled - this->mTestDraftTokenIdsInit = { - {4, 1, 2}, // - {5, 0, 3}, // - {4, 1, 2}, // - {5, 0, 3}, // - {4, 1, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {5}, {4}, {5}, {4}, {4, 5}, // step 0 - {1}, {0}, {1}, {0}, {1}, {0}, // step 1 - {2}, {3}, {2}, {3}, {2}, {0}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopKReject) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 2}, // accept, accept, accept, sampled - {4, 3, 4}, // accept, reject, 0, 0 - {4, 3, 4}, // accept, reject, 0, 0 - {2, 3, 4}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4, 5, 6, 7}, {4, 5, 6, 7}, // step 0 - {0}, {0}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0, 1}, {0, 1}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopK1TopP0) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopK1TopP0Reject) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 3}, // accept, accept, reject, 0 - {4, 1, 2}, // accept, reject, 0, 0 - {4, 1, 2}, // accept, reject, 0, 0 - {5, 0, 2}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsBatchTopK) -{ - std::vector topKs = {1, 1, 2, 2, 4, 4}; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topKs}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {}, // no draft tokens, token will be sampled - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0, 1}, {0, 1}, {0, 1, 2, 3}, {0, 1, 2, 3}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsBatchTopKReject) -{ - std::vector topKs = {1, 1, 2, 2, 4, 4}; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topKs}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {5, 0, 2}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - {4, 0, 4}, // accept, accept, reject, 0 - {4, 0, 2}, // accept, accept, accept, sampled - {4, 5, 2}, // accept, reject, 0, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4, 5, 6, 7}, {4, 5}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0, 1, 2, 3}, // step 1 - {2}, {0}, {0}, {2, 3}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0, 1, 2, 3}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopP) -{ - // Skip topK decode - float topP = 0.3; - TestSamplingParams params; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopPReject) -{ - // Skip topK decode - float topP = 0.3; - TestSamplingParams params; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 2}, // accept, accept, reject, 0 - {4, 1, 3}, // accept, reject, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopKTopP) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {4, 0, 2}, // - {}, // no draft tokens, token will be sampled - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsTopKTopPReject) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, // accept, accept, accept, sampled - {4, 0, 3}, // accept, accept, reject, 0 - {4, 3, 2}, // accept, reject, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - {7, 0, 2}, // reject, 0, 0, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsBatchTopKBatchTopP) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {6, 2, 4}, - {4, 0, 2}, - {4, 0, 2}, - {5, 1, 3}, - {4, 0, 2}, - {4, 0, 2}, - }; - - std::vector> expectedOutputIds{ - // batch - {6}, {4}, {4}, {5}, {4}, {4}, // step 0 - {2}, {0}, {0}, {1}, {0}, {0}, // step 1 - {4}, {2}, {2}, {3}, {2}, {2}, // step 2 - {0, 1, 2}, {0}, {0}, {0, 1}, {0}, {0, 1}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsBatchTopKBatchTopPReject) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {6, 3, 4}, // accept, reject, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - {6, 0, 2}, // reject, 0, 0, 0 - {7, 1, 3}, // reject, 0, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - {4, 2, 2}, // accept, reject, 0, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {6}, {4}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2}, {0}, {0}, {0}, {0}, {0, 1}, // step 1 - {0}, {2}, {0}, {0}, {2}, {0}, // step 2 - {0, 1, 2}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsBatchTopK0BatchTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsAccept; - - this->mTestDraftTokenIdsInit = { - {7, 3, 5}, - {5, 1, 3}, - {5, 1, 3}, - {5, 1, 3}, - {4, 0, 2}, - {4, 0, 2}, - }; - - std::vector> expectedOutputIds{ - // batch - {7}, {5}, {5}, {5}, {4}, {4}, // step 0 - {3}, {1}, {1}, {1}, {0}, {0}, // step 1 - {5}, {3}, {3}, {3}, {2}, {2}, // step 2 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByLogitsBatchTopK0BatchTopPReject) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = true; - params.isAirTopPExternalDraftTokensLayer = true; - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1) - this->mTestDraftLogitsInit = &this->mTestDraftLogitsReject; - - this->mTestDraftTokenIdsInit = { - {7, 4, 5}, // accept, reject, 0, 0 - {5, 5, 3}, // accept, reject, 0, 0 - {6, 1, 3}, // reject, 0, 0, 0 - {5, 1, 3}, // accept, accept, accept, sampled - {4, 2, 2}, // accept, reject, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {7}, {5}, {4, 5}, {5}, {4}, {4}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0}, {1}, {0}, {0}, // step 1 - {0}, {0}, {0}, {3}, {0}, {2}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByTokenIdsBatchTopKBatchTopP) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - params.isAirTopPExternalDraftTokensLayer = true; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {4, 0, 2}, - {4, 0, 2}, - {4, 0, 2}, - {4, 0, 2}, - {4, 0, 2}, - {4, 0, 2}, - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {0, 2, 3, 4}, {2}, {2}, {0, 2, 3}, {2}, {2}, // step 2 - {0, 1, 2}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByTokenIdsBatchTopKBatchTopPReject) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - params.isAirTopPExternalDraftTokensLayer = true; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {4, 3, 2}, // accept, reject, 0, 0 - {5, 0, 2}, // reject, 0, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - {6, 0, 2}, // reject, 0, 0, 0 - {4, 1, 2}, // accept, reject, 0, 0 - {4, 1, 2}, // accept, reject, 0, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {0}, {2}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByTokenIdsBatchTopK0BatchTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - params.isAirTopPExternalDraftTokensLayer = true; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {7, 3, 5}, - {5, 1, 3}, - {5, 1, 3}, - {4, 0, 2}, - {4, 0, 2}, - {4, 0, 2}, - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0}, // step 1 - {0, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 2, 3}, {0, 2, 3}, {2}, {2}, // step 2 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, AirTopPAcceptByTokenIdsBatchTopK0BatchTopPReject) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - std::vector topPs = {1.0, 1.0, 0.5, 0.5, 0.3, 0.3}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - params.isAirTopPExternalDraftTokensLayer = true; - - // accept by token ids result may different for different seeds - // therefore there are more possible paths in expectedOutputIds - this->mTestDraftTokenIdsInit = { - {7, 4, 5}, // accept, reject, 0, 0 - {5, 1, 6}, // accept/reject, accept/reject, reject, 0 - {6, 1, 3}, // reject, 0, 0, 0 - {4, 0, 2}, // accept/reject, accept/reject, accept, sampled - {4, 1, 2}, // accept, reject, 0, 0 - {4, 0, 3}, // accept, accept, reject, 0 - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6, 7}, {4, 5, 6, 7}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1, 2, 3}, {0, 1, 2, 3}, {0}, {0, 1}, {0}, {0}, // step 1 - {0}, {0, 2, 3, 4, 5}, {0}, {0, 2, 3}, {0}, {2}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(ExternalDraftTokensLayerTest, BatchTopKBatchTopP) -{ - std::vector topKs = {3, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - params.isExternalDraftTokensLayerTest = true; - params.useDraftLogits = false; - - this->mTestDraftTokenIdsInit = { - {}, - {}, - {}, - {}, - {}, - {}, - }; - - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0}, {0, 1}, {0}, {0}, {0}, {0}, // step 1 - {0}, {0, 1}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0, 1}, {0}, {0}, {0}, {0}, // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -} // namespace diff --git a/cpp/tests/unit_tests/layers/layerUtilsTest.cpp b/cpp/tests/unit_tests/layers/layerUtilsTest.cpp deleted file mode 100644 index 468693214ece..000000000000 --- a/cpp/tests/unit_tests/layers/layerUtilsTest.cpp +++ /dev/null @@ -1,255 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include -#include -#include -#include - -namespace tensorrt_llm::tests::layers -{ - -using namespace tensorrt_llm::runtime; - -class MineFieldAllocator : public std::pmr::memory_resource -{ - static constexpr std::size_t kMinPadding = 256; - - void* do_allocate(std::size_t bytes, std::size_t alignment) override - { - alignment = std::max(alignment, sizeof(uint64_t)); - bytes = common::roundUp(bytes, alignment); - auto const padding = common::roundUp(kMinPadding, alignment); - - auto const allocSize = bytes + 2 * padding; - void* p = std::pmr::new_delete_resource()->allocate(allocSize, alignment); - - std::generate_n(static_cast(p), allocSize / sizeof(uint64_t), - [engine = std::mt19937_64(reinterpret_cast(p))]() mutable { return engine(); }); - - return static_cast(p) + padding; - } - - void do_deallocate(void* p, std::size_t bytes, std::size_t alignment) override - { - auto allocAlignment = std::max(alignment, sizeof(uint64_t)); - auto allocBytes = common::roundUp(bytes, allocAlignment); - auto const padding = common::roundUp(kMinPadding, alignment); - void* allocP = static_cast(p) - padding; - auto const allocSize = allocBytes + 2 * padding; - - auto engine = std::mt19937_64(reinterpret_cast(allocP)); - auto* verifyP = static_cast(allocP); - for (size_t i = 0; i < padding / sizeof(uint64_t); ++i) - { - ASSERT_EQ(verifyP[i], engine()); - } - - engine.discard(allocBytes / sizeof(uint64_t)); - - verifyP += (allocBytes + padding) / sizeof(uint64_t); - for (size_t i = 0; i < padding / sizeof(uint64_t); ++i) - { - ASSERT_EQ(verifyP[i], engine()); - } - - std::pmr::new_delete_resource()->deallocate(allocP, allocSize, allocAlignment); - } - - bool do_is_equal(std::pmr::memory_resource const& other) const noexcept override - { - return *this == other; - } -}; - -MineFieldAllocator mineFieldAllocator{}; - -template -class CopyToWorkspaceFixture : public testing::Test -{ -public: - CopyToWorkspaceFixture() - { - bufferManager = std::make_unique(std::make_shared()); - } - - void SetUp() override - { - std::pmr::set_default_resource(&mineFieldAllocator); - } - - void fillData(std::pmr::vector& vec) - { - if constexpr (std::is_pointer_v) - { - std::iota(vec.begin(), vec.end(), static_cast(nullptr)); - } - else - { - std::iota(vec.begin(), vec.end(), 0); - } - } - - std::unique_ptr bufferManager; - static T value_; -}; - -using DataTypes = ::testing::Types; - -TYPED_TEST_SUITE(CopyToWorkspaceFixture, DataTypes); - -TYPED_TEST(CopyToWorkspaceFixture, DataTooLarge_Throws) -{ - using dataType = decltype(this->value_); - constexpr size_t numElements = 1024; - auto const data = std::pmr::vector(numElements); - auto const workspaceSizeInBytes = numElements * sizeof(dataType) / 2; - IBuffer::SharedPtr workspace = this->bufferManager->gpu(workspaceSizeInBytes); - ASSERT_THROW(DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace), common::TllmException); -} - -TYPED_TEST(CopyToWorkspaceFixture, DataMuchTooLarge_Throws) -{ - using dataType = decltype(this->value_); - constexpr size_t numElements = 1 << 15; - auto const data = std::pmr::vector(numElements); - auto const workspaceSizeInBytes = numElements * sizeof(dataType) / 1000; - IBuffer::SharedPtr workspace = this->bufferManager->gpu(workspaceSizeInBytes); - ASSERT_THROW(DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace), common::TllmException); -} - -TYPED_TEST(CopyToWorkspaceFixture, DataFitsExactly_Succeeds) -{ - using dataType = decltype(this->value_); - constexpr size_t numElements = 2048; - auto data = std::pmr::vector(numElements); - this->fillData(data); - auto const workspaceSizeInBytes = numElements * sizeof(dataType); - IBuffer::SharedPtr workspace = this->bufferManager->gpu(workspaceSizeInBytes); - DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - // Copy back and check data integrity. - auto dataCopy = std::pmr::vector(numElements); - auto const dataSizeInBytes = numElements * sizeof(dataType); - this->bufferManager->copy(*IBuffer::slice(workspace, 0, dataSizeInBytes), dataCopy.data(), MemoryType::kCPU); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - for (auto i = 0; i < numElements; i++) - { - ASSERT_EQ(dataCopy[i], data[i]); - } -} - -TYPED_TEST(CopyToWorkspaceFixture, DataSmallerThanWorkspace_Succeeds) -{ - using dataType = decltype(this->value_); - constexpr size_t numElements = 2048; - auto data = std::pmr::vector(numElements); - this->fillData(data); - auto const workspaceSizeInBytes = numElements * sizeof(dataType) * 4; - IBuffer::SharedPtr workspace = this->bufferManager->gpu(workspaceSizeInBytes); - DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - // Copy back and check data integrity. - auto dataCopy = std::pmr::vector(numElements); - auto const dataSizeInBytes = numElements * sizeof(dataType); - this->bufferManager->copy(*IBuffer::slice(workspace, 0, dataSizeInBytes), dataCopy.data(), MemoryType::kCPU); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - for (auto i = 0; i < numElements; i++) - { - ASSERT_EQ(dataCopy[i], data[i]); - } -} - -TYPED_TEST(CopyToWorkspaceFixture, DataMuchSmallerThanWorkspace_Succeeds) -{ - using dataType = decltype(this->value_); - constexpr size_t numElements = 2048; - auto data = std::pmr::vector(numElements); - this->fillData(data); - auto const workspaceSizeInBytes = numElements * sizeof(dataType) * 1000; - IBuffer::SharedPtr workspace = this->bufferManager->gpu(workspaceSizeInBytes); - DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - // Copy back and check data integrity. - auto dataCopy = std::pmr::vector(numElements); - auto const dataSizeInBytes = numElements * sizeof(dataType); - this->bufferManager->copy(*IBuffer::slice(workspace, 0, dataSizeInBytes), dataCopy.data(), MemoryType::kCPU); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - for (auto i = 0; i < numElements; i++) - { - ASSERT_EQ(dataCopy[i], data[i]); - } -} - -TYPED_TEST(CopyToWorkspaceFixture, TypedWorkspaceBuffer_Succeeds) -{ - using dataType = decltype(this->value_); - if constexpr (std::is_same_v) - { - // There's no TRTDataType - } - else - { - constexpr size_t numElements = 2048; - auto data = std::pmr::vector(numElements); - this->fillData(data); - IBuffer::SharedPtr workspace = this->bufferManager->gpu(numElements, TRTDataType::value); - DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - // Copy back and check data integrity. - auto dataCopy = std::pmr::vector(numElements); - this->bufferManager->copy(*IBuffer::slice(workspace, 0, numElements), dataCopy.data(), MemoryType::kCPU); - sync_check_cuda_error(this->bufferManager->getStream().get()); - - for (auto i = 0; i < numElements; i++) - { - ASSERT_EQ(dataCopy[i], data[i]); - } - } -} - -TYPED_TEST(CopyToWorkspaceFixture, MismatchBufferType_Throws) -{ - using dataType = decltype(this->value_); - if constexpr (sizeof(dataType) == 1 || std::is_same_v) - { - // Allow copy mismatch type into int8_t workspace buffer - // There's no TRTDataType - } - else - { - constexpr size_t numElements = 2048; - using differentType = std::pair; - auto data = std::pmr::vector(numElements); - IBuffer::SharedPtr workspace = this->bufferManager->gpu(numElements, TRTDataType::value); - ASSERT_THROW( - DecodingLayerWorkspace::copyToWorkspace(*this->bufferManager, data, workspace), common::TllmException); - } -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp b/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp deleted file mode 100644 index 7f4791a85ba4..000000000000 --- a/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp +++ /dev/null @@ -1,264 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include -#include - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/layers/lookaheadAlgorithm.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/lookaheadModule.h" -#include "tests/unit_tests/layers/randomLlm.h" - -namespace tensorrt_llm::tests::layers -{ -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using TensorPtr = runtime::ITensor::SharedPtr; - -class LookaheadAlgorithmTest - : public ::testing::TestWithParam, std::tuple, std::tuple>> -{ -}; - -bool verifyAcceptOffsets(TensorPtr output, TensorPtr accepted, TensorPtr acceptedOffsets) -{ - BufferRange outputRange(*output); - BufferRange acceptedRange(*accepted); - BufferRange offsetsRange(*acceptedOffsets); - bool result = true; - for (SizeType32 i = 1; i < acceptedRange.size(); i++) - { - result &= outputRange[offsetsRange[i - 1] + 1] == acceptedRange[i]; - } - return result; -} - -TEST_P(LookaheadAlgorithmTest, predict) -{ - srand(42); - auto [Ww, Nn, Gg] = GetParam(); - auto [W, w] = Ww; - auto [N, n] = Nn; - auto [G, g] = Gg; - - if (!executor::LookaheadDecodingConfig::isLegal(W, N, G) || !executor::LookaheadDecodingConfig::isLegal(w, n, g)) - { - TLLM_LOG_DEBUG("Just Pass for illegal parameter combination"); - GTEST_SKIP() << "Algorithm does not support these parameters WNG=(" << W << ", " << N << ", " << G << "), wng=(" - << w << ", " << n << ", " << g; - } - TLLM_LOG_DEBUG("Test Parameters: WNG=(%d, %d, %d), wng=(%d, %d, %d)", W, N, G, w, n, g); - - auto ascii = std::make_shared(); - - std::string oracle( - "The following example uses the following lambda-expression to increment all of the elements of a vector and " - "then uses an overloaded operator() in a function object (a.k.a., \"functor\") to compute their sum. Note that " - "to compute the sum, it is recommended to use the dedicated algorithm std::accumulate.&"); - LookaheadRandomLlm llm(ascii, oracle); - - auto prompt = initTensor(std::string(oracle.substr(0, 20))); - BufferRange promptRange(*prompt); - auto promptLen = ITensor::volume(prompt->getShape()); - - auto maxSeqLen = 1024; - SizeType32 maxTokensPerStep, maxDraftLen; - SizeType32 maxDraftLenRuntime; - std::tie(maxTokensPerStep, std::ignore, maxDraftLen, std::ignore) - = executor::LookaheadDecodingConfig(W, N, G).calculateSpeculativeResource(); - std::tie(std::ignore, std::ignore, maxDraftLenRuntime, std::ignore) - = executor::LookaheadDecodingConfig(w, n, g).calculateSpeculativeResource(); - auto shape = ITensor::makeShape({maxTokensPerStep}); - auto shape2d = ITensor::makeShape({maxTokensPerStep, maxTokensPerStep}); - auto shapeSingle = ITensor::makeShape({1}); - TensorPtr posidMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr attentionMaskMax = BufferManager::cpu(shape2d, tensorrt_llm::DataType::kBOOL); - TensorPtr inputLengthPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); - auto& inputLength(*BufferRange(*inputLengthPtr).begin()); - - TensorPtr outputMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr endIdPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); - auto& endId(*BufferRange(*endIdPtr).begin()); - endId = ascii->getEndToken(); - - TensorPtr acceptedMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr acceptedOffsetsMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr acceptedLengthPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); - auto& acceptedLength(*BufferRange(*acceptedLengthPtr).begin()); - - TensorPtr sequence - = BufferManager::cpu(ITensor::makeShape({maxSeqLen + maxDraftLen}), tensorrt_llm::DataType::kINT32); - BufferRange sequenceRange(*sequence); - TensorPtr sequenceLengthPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); - auto& sequenceLength(*bufferCast(*sequenceLengthPtr)); - - std::copy(promptRange.begin(), promptRange.end(), sequenceRange.begin()); - sequenceLength = promptLen; - - sequenceRange[sequenceLength] = oracle[promptLen]; // from context phase. - sequenceLength += 1; - - PRINT_TOKENS(sequence); - - tensorrt_llm::layers::LookaheadAlgorithm algo(W, N, G); - algo.setup(ITensor::slice(sequence, 0, sequenceLength), w, n, g, 42); - - SizeType32 seqLen = oracle.size(); - std::vector histogram(N + 1); - - for (; sequenceLength < seqLen;) - { - TLLM_LOG_DEBUG("\noracle[%d] = '%c'", sequenceLength - 1, static_cast(sequenceRange[sequenceLength - 1])); - bufferCast(*posidMax)[0] = sequenceLength - 1; - BufferLocation amaskLocation(*attentionMaskMax); - for (auto& item : amaskLocation) - { - item = false; - } - for (SizeType32 i = 0; i < maxTokensPerStep; i++) - { - amaskLocation.at(i, 0) = true; - } - - algo.prepare( // - ITensor::slice(sequence, sequenceLength, maxDraftLenRuntime), // - ITensor::slice(posidMax, 1, maxDraftLenRuntime), // - inputLengthPtr, // - attentionMaskMax, 1, // - sequenceLengthPtr, // - ITensor::slice(sequence, sequenceLength - 1, 1)); - - TensorPtr input = ITensor::slice(sequence, sequenceLength - 1, inputLength + 1); - TensorPtr posid = ITensor::slice(posidMax, 0, inputLength + 1); - TensorPtr amask = ITensor::slice(attentionMaskMax, 0, inputLength + 1); - - PRINT_TOKENS(input); - PRINT_VALUES(posid); - PRINT_VALUES(amask); - - TensorPtr output = ITensor::slice(outputMax, 0, inputLength + 1); - llm.foretell(output, input, posid, amask); - PRINT_TOKENS(output); - - // algo.update(acceptedMax, acceptedOffsetsMax, acceptedLengthPtr, output, endIdPtr); - algo.update( - ITensor::slice(sequence, sequenceLength, n), acceptedOffsetsMax, acceptedLengthPtr, output, endIdPtr); - - TensorPtr accepted = ITensor::slice(sequence, sequenceLength, acceptedLength); - TensorPtr acceptedOffsets = ITensor::slice(acceptedOffsetsMax, 0, acceptedLength); - - TLLM_CHECK(acceptedLength <= N); - histogram[acceptedLength] += 1; - PRINT_TOKENS(accepted); - PRINT_VALUES(acceptedOffsets); - - EXPECT_TRUE(verifyAcceptOffsets(output, accepted, acceptedOffsets)); - EXPECT_TRUE(llm.verify(sequenceLength, accepted)); - - sequenceLength += acceptedLength; - - TLLM_LOG_DEBUG("result: '%s'", D(ITensor::slice(sequence, 0, sequenceLength)).string().c_str()); - } - EXPECT_EQ(sequenceLength, seqLen); - - TensorPtr hist = ITensor::wrap(histogram, ITensor::makeShape({N + 1})); - TLLM_LOG_DEBUG("Lookahead acceptance histogram: %s", D(hist).values().c_str()); -} - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTest, LookaheadAlgorithmTest, - testing::Combine( // - testing::Values(std::make_tuple(1, 1), std::make_tuple(3, 3), std::make_tuple(5, 5), std::make_tuple(7, 7), - std::make_tuple(2, 1), std::make_tuple(3, 2), std::make_tuple(5, 3), std::make_tuple(7, 4)), - testing::Values(std::make_tuple(1, 1), std::make_tuple(3, 3), std::make_tuple(5, 5), std::make_tuple(7, 7), - std::make_tuple(2, 1), std::make_tuple(3, 2), std::make_tuple(5, 3), std::make_tuple(7, 4)), - testing::Values(std::make_tuple(0, 0), std::make_tuple(3, 3), std::make_tuple(5, 5), std::make_tuple(7, 7), - std::make_tuple(1, 0), std::make_tuple(3, 2), std::make_tuple(5, 3), std::make_tuple(7, 4)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSingleMax, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(5, 5)), testing::Values(std::make_tuple(5, 5)), - testing::Values(std::make_tuple(5, 5)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSingleDynamic, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(1, 1)), testing::Values(std::make_tuple(2, 1)), - testing::Values(std::make_tuple(1, 0)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSmallest_110, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(1, 1)), testing::Values(std::make_tuple(1, 1)), - testing::Values(std::make_tuple(0, 0)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSmall_120, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(1, 1)), testing::Values(std::make_tuple(2, 2)), - testing::Values(std::make_tuple(0, 0)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSmall_220, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(2, 2)), testing::Values(std::make_tuple(2, 2)), - testing::Values(std::make_tuple(0, 0)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSmall_121, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(1, 1)), testing::Values(std::make_tuple(2, 2)), - testing::Values(std::make_tuple(1, 1)))); - -INSTANTIATE_TEST_CASE_P(CombineLookaheadAlgorithmTestSmall_222, LookaheadAlgorithmTest, - testing::Combine(testing::Values(std::make_tuple(2, 2)), testing::Values(std::make_tuple(2, 2)), - testing::Values(std::make_tuple(2, 2)))); - -TEST(LookaheadAlgorithmTest, treeEncodeTest) -{ - auto testWithData = [](TensorPtr inputTokens, TensorPtr inputPosIds, SizeType32 lastPosId, SizeType32 gold_len) - { - auto shape = inputTokens->getShape(); - auto shape2d = ITensor::makeShape({shape.d[0], shape.d[0]}); - - TensorPtr inputMasks = BufferManager::cpu(shape2d, tensorrt_llm::DataType::kBOOL); - LookaheadAlgorithm::posIdsToMask(inputMasks, inputPosIds); - - TensorPtr outputTokens = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr outputPosIds = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr encodeMap = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); - TensorPtr outputMasks = BufferManager::cpu(shape2d, tensorrt_llm::DataType::kBOOL); - - // auto len = LookaheadAlgorithm::treeEncode(outputTokens, outputPosIds, outputMasks, inputTokens, inputPosIds, - // inputMasks, '$', 9); - auto len = LookaheadAlgorithm::treeEncode(inputTokens, inputPosIds, inputMasks, encodeMap); - TLLM_LOG_DEBUG("len = %d", len); - - EXPECT_EQ(len, gold_len); - }; - - testWithData( // - initTensor(std::string("01234512345")), // - initTensor({10, 11, 12, 13, 14, 15, 11, 12, 13, 14, 15}), // - 9, 6); - - testWithData( // - initTensor(std::string("01234512abc")), // - initTensor({10, 11, 12, 13, 14, 15, 11, 12, 13, 14, 15}), // - 9, 9); - - testWithData( // - initTensor(std::string("01234512abc2aBCD")), // - initTensor({10, 11, 12, 13, 14, 15, 11, 12, 13, 14, 15, 12, 13, 14, 15, 16}), // - 9, 12); - - testWithData(initTensor(std::string("wmplhi folxamp")), - initTensor({21, 22, 23, 24, 25, 26, 27, 21, 22, 23, 24, 21, 22, 23, 24}), 20, 15); -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp b/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp deleted file mode 100644 index 917f6dbdca55..000000000000 --- a/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp +++ /dev/null @@ -1,873 +0,0 @@ -/* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/layers/decodingParams.h" -#include "tensorrt_llm/layers/lookaheadDecodingLayer.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/lookaheadModule.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tests/unit_tests/layers/randomLlm.h" - -namespace tensorrt_llm::tests::layers -{ -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; - -namespace trk = tensorrt_llm::runtime::kernels; - -using TensorPtr = runtime::ITensor::SharedPtr; -using TensorConstPtr = runtime::ITensor::SharedConstPtr; - -struct TestParam -{ - SizeType32 maxBatchSize; - - enum BatchType - { - SINGLE_ONCE, - SINGLE_TWICE, - DYNAMIC - } batchType; - - SizeType32 maxW; - SizeType32 w; - SizeType32 maxN; - SizeType32 n; - SizeType32 maxG; - SizeType32 g; -}; - -class BatchSlotsManager -{ -public: - BatchSlotsManager(SizeType32 maxBatchSize, SizeType32 cases) - : mMaxBatchSize(maxBatchSize) - , mCases(cases) - { - } - - virtual std::vector alloc(void) = 0; - virtual void free(SizeType32 id) = 0; - - bool finished() - { - return mCases == 0; - } - -protected: - SizeType32 quota(void) - { - return mCases - mRunning; - } - - void consume(SizeType32 cases) - { - TLLM_CHECK(cases >= 0); - TLLM_CHECK_DEBUG_WITH_INFO(cases <= mCases, "cases=%d, mCases=%d", cases, mCases); - mRunning -= cases; - mCases -= cases; - } - -protected: - SizeType32 mMaxBatchSize{0}; - SizeType32 mCases{0}; - SizeType32 mRunning{0}; -}; - -class SingleBatchSlotsManager : public BatchSlotsManager -{ -public: - SingleBatchSlotsManager(SizeType32 maxBatchSize, SizeType32 cases, SizeType32 id) - : BatchSlotsManager(maxBatchSize, cases) - , mId(id) - { - TLLM_CHECK(id < maxBatchSize); - } - - virtual std::vector alloc(void) - { - if (mState == FREE && quota() > 0) - { - mState = BUSY; - mRunning += 1; - return std::vector({mId}); - } - else - { - return std::vector(); - } - } - - virtual void free(SizeType32 id) - { - TLLM_CHECK(id == mId); - mState = FREE; - consume(1); - } - -private: - enum - { - FREE, - BUSY - } mState{FREE}; - - SizeType32 mId; -}; - -class DynamicBatchSlotsManager : public BatchSlotsManager -{ -public: - DynamicBatchSlotsManager(SizeType32 maxBatchSize, SizeType32 cases) - : BatchSlotsManager(maxBatchSize, cases) - { - for (SizeType32 bi = 0; bi * 3 + 2 < maxBatchSize; bi++) - { - mFreeList.push(bi * 3 + 1); - mFreeList.push(bi * 3 + 2); - mFreeList.push(bi * 3); - } - } - - virtual std::vector alloc() - { - SizeType32 waterline = mMaxBatchSize / 4; - SizeType32 plan = mBusySet.size() < waterline ? rand() % (mMaxBatchSize / 4) : 0; - SizeType32 num = std::min(plan, quota()); - std::vector result; - for (SizeType32 i = 0; i < num && !mFreeList.empty(); i++) - { - SizeType32 id = mFreeList.front(); - result.push_back(id); - mBusySet.insert(id); - mFreeList.pop(); - } - mRunning += result.size(); - return result; - } - - virtual void free(SizeType32 id) - { - auto search = mBusySet.find(id); - TLLM_CHECK(search != mBusySet.end()); - mBusySet.erase(search); - mFreeList.push(id); - consume(1); - } - -private: - std::queue mFreeList; - std::set mBusySet; -}; - -class LookaheadDecodingLayerTest : public testing::Test -{ -public: - void SetUp() override; - void TearDown() override; - void runTest(TestParam const& param); - -private: - void allocateBuffers(); - - void setupBuffers(); - - void newRequests(std::vector requestIds); - - void manageBatch(); - - void llmForward(); - - void decodeForward(); - - void verifyDecode(); - -protected: - std::shared_ptr mBufferManager; - std::shared_ptr mStream; - - struct cudaDeviceProp mDeviceProp; - - TensorPtr mAlgoConfigBatch; - - TensorPtr mOutputIds; - TensorPtr mSequenceLengths; - TensorPtr mProbs; - TensorPtr mEndIds; - TensorPtr mTokensPerStep; - TensorPtr mGoldenSampledTokens; - TensorPtr mBatchSlots; - TensorPtr mBatchSlotsMax; - - TensorPtr mNewTokens; - TensorPtr mNumNewTokens; - TensorPtr mNumNewTokensCumSum; - TensorPtr mPathsOffsets; - TensorPtr mDraftLengths; - TensorPtr mPrevDraftLengths; - TensorPtr mDraftTokens; - TensorPtr mPackedMasks; - TensorPtr mPackedMasksBool; - TensorPtr mGenerationLengths; - TensorPtr mPositionOffsets; - TensorPtr mPositionIds; - TensorPtr mAttentionPackedMask; - - TensorPtr mInputTokensBatch; - TensorPtr mPositionIdsBatch; - - int32_t mMaxTopK = 1; - static constexpr int32_t mMaxSeqLen = 512; - float mMaxTopP = 1.0; - std::shared_ptr mAscii; - std::vector mOracle; - std::vector mPrompt; - std::vector> mLlm; - std::shared_ptr> mDecoder; - std::shared_ptr mDecodingWorkspace; - SizeType32 mVocabSize; - SizeType32 mMaxTokensPerStep; - TestParam mTestParam; - std::shared_ptr mBatchSlotsManager; - std::vector mScoreBoard; - std::vector mHistogram; - std::list mReports; -}; - -void LookaheadDecodingLayerTest::SetUp() -{ - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); - - int32_t device = 0; - cudaGetDevice(&device); - cudaGetDeviceProperties(&mDeviceProp, device); - - mAscii = std::make_shared(); - mVocabSize = mAscii->getVocabSize(); -} - -void LookaheadDecodingLayerTest::TearDown() {} - -void LookaheadDecodingLayerTest::allocateBuffers() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const maxBatchSize = mTestParam.maxBatchSize; - auto const vocabSize = mAscii->getVocabSize(); - auto const maxBeamSize = 1; - - SizeType32 maxNumNewTokens, maxDraftLen, maxAcceptedDraftLen; - std::tie(mMaxTokensPerStep, maxNumNewTokens, maxDraftLen, maxAcceptedDraftLen) - = executor::LookaheadDecodingConfig(mTestParam.maxW, mTestParam.maxN, mTestParam.maxG) - .calculateSpeculativeResource(); - // mMaxTokensPerStep = maxTokensPerStep; - - auto const vocabSizePadded = vocabSize; - auto const maxNumHeads = 1; - std::ostringstream buf; - - std::vector text({// - std::string("To be, or not to be: that is the question. " - "To Be, Or Not To Be: That Is The Question.&"), - std::string("Be not afraid of greatness. Some are born great, some achieve greatness, and others have " - "greatness thrust upon them. " - "Be Not Afraid Of Greatness. Some Are Born Great, Some Achieve Greatness, And Others Have " - "Greatness Thrust Upon Them.&"), - std::string("Sweet are the uses of adversity which, like the toad, ugly and venomous, wears yet a precious " - "jewel in his head. " - "Sweet Are the Uses Of Adversity Which, Like The Toad, Ugly And Venomous, Wears Yet A Precious " - "Jewel In His Head.&"), - std::string("Talking isn't doing. It is a kind of good deed to say well; and yet words are not deeds. " - "Talking Isn't Doing. It Is A Kind Of Good Deed To Say Well; And Yet Words Are Not Deeds.&"), - std::string( - "Reputation is an idle and most false imposition; oft got without merit, and lost without deserving. " - "Reputation Is An Idle And Most False Imposition; Oft Got Without Merit, And Lost Without Deserving.&")}); - - mOracle.resize(maxBatchSize); - mLlm.resize(maxBatchSize); - mPrompt.resize(maxBatchSize); - mScoreBoard.resize(maxBatchSize); - mHistogram.resize(maxBatchSize); - for (SizeType32 gbi = 0; gbi < maxBatchSize; gbi++) - { - mOracle[gbi] = text[rand() % text.size()]; - mLlm[gbi] = std::make_shared(mAscii, mOracle[gbi], gbi); - - mScoreBoard[gbi] = std::ostringstream(); - mHistogram[gbi] = BufferManager::cpu(ITensor::makeShape({mTestParam.n + 1}), tensorrt_llm::DataType::kINT32); - } - switch (mTestParam.batchType) - { - case TestParam::SINGLE_ONCE: - mBatchSlotsManager = std::make_shared(maxBatchSize, 1, 1); - break; - case TestParam::SINGLE_TWICE: - mBatchSlotsManager = std::make_shared(maxBatchSize, 2, 1); - break; - case TestParam::DYNAMIC: - mBatchSlotsManager = std::make_shared(maxBatchSize, maxBatchSize * 2); - break; - } - - auto lookaheadModule = std::make_shared(mTestParam.maxN, mMaxTokensPerStep - 1); - - lookaheadModule->setExecutionConfig( - executor::LookaheadDecodingConfig(mTestParam.maxW, mTestParam.maxN, mTestParam.maxG)); - auto const decodingDomain - = tensorrt_llm::layers::DecoderDomain(maxBatchSize, 1, vocabSize, vocabSizePadded, lookaheadModule); - - mDecoder = std::make_shared>(decodingDomain, mBufferManager); - - TLLM_LOG_DEBUG("decoder ok"); - - auto maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - - mAlgoConfigBatch = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, 3}), tensorrt_llm::DataType::kINT32); - - mEndIds = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mTokensPerStep = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - - mOutputIds - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, maxBeamSize, mMaxSeqLen + mMaxTokensPerStep}), - tensorrt_llm::DataType::kINT32); - mSequenceLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - - mProbs = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep, vocabSize}), tensorrt_llm::DataType::kFLOAT); - - mGoldenSampledTokens - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mInputTokensBatch = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mPositionIdsBatch = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); - - mNewTokens = BufferManager::pinnedPool( - ITensor::makeShape({mMaxTokensPerStep, maxBatchSize, 1}), tensorrt_llm::DataType::kINT32); - mNumNewTokens = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mDraftLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mPrevDraftLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mDraftTokens - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, maxDraftLen}), tensorrt_llm::DataType::kINT32); - auto packedMaskShape = ITensor::makeShape( - {maxBatchSize, mMaxTokensPerStep, static_cast(common::divUp(mMaxTokensPerStep, 32))}); - mPackedMasks = BufferManager::pinnedPool(packedMaskShape, tensorrt_llm::DataType::kINT32); - mPackedMasksBool = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep, mMaxTokensPerStep}), tensorrt_llm::DataType::kBOOL); - mNumNewTokensCumSum - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize + 1}), tensorrt_llm::DataType::kINT32); - mPathsOffsets = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); - mGenerationLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mPositionOffsets = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mPositionIds = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); - mAttentionPackedMask = BufferManager::pinnedPool(packedMaskShape, tensorrt_llm::DataType::kINT32); - - mBatchSlotsMax = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - - auto const batchSize = 0; - auto batchShape1D = ITensor::makeShape({batchSize}); - auto batchShape2D = ITensor::makeShape({batchSize, mMaxTokensPerStep}); - - mBatchSlots = ITensor::slice(mBatchSlotsMax, 0, batchSize); - - trk::invokeFill(*mEndIds, mAscii->getEndToken(), *mStream); - trk::invokeFill(*mOutputIds, int32_t{0}, *mStream); - trk::invokeFill(*mSequenceLengths, int32_t{0}, *mStream); - trk::invokeFill(*mTokensPerStep, mMaxTokensPerStep, *mStream); - mDecodingWorkspace = std::make_unique( - mBufferManager, decodingDomain, TRTDataType::value, mDecoder->getWorkspaceSize()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadDecodingLayerTest::setupBuffers() {} - -void LookaheadDecodingLayerTest::newRequests(std::vector requestIds) -{ - TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); - auto const requestSize = requestIds.size(); - - auto const beamSize = 1; - SizeType32 vocabSize = mAscii->getVocabSize(); - - //////////////////////////////// - for (auto gbi : requestIds) - { - auto len = 5 + rand() % 10; - auto prompt = mOracle[gbi].substr(0, len); - - TokenIdType contextToken = mOracle[gbi][len]; - SizeType32 contextLen = len + 1; - - BufferRange outputRange(*ITensor::at(mOutputIds, {gbi, 0})); - for (auto& v : outputRange) - { - v = 0; - } - std::copy(prompt.begin(), prompt.end(), outputRange.begin()); - outputRange[len] = contextToken; - BufferLocation(*mSequenceLengths).at(gbi) = len + 1; - BufferLocation(*mDraftLengths).at(gbi) = 0; - BufferLocation(*mNumNewTokens).at(gbi) = 0; - - mPrompt[gbi] = ITensor::slice(mOutputIds, {gbi, 0, 0}, len + 1); - - for (auto& v : BufferRange(*mHistogram[gbi])) - { - v = 0; - } - mScoreBoard[gbi] << "request id=[" << gbi << "] starts. prompt len=[" << len << "]."; - } - - TLLM_LOG_DEBUG("batch slots"); - //////////////////////////////// - auto batchSize = ITensor::volume(mBatchSlots->getShape()); - BufferRange batchSlotMaxRange(*mBatchSlotsMax); - std::copy(requestIds.begin(), requestIds.end(), batchSlotMaxRange.begin() + batchSize); - - //////////////////////////////// - auto setupParams = std::make_shared(); - setupParams->prompt.resize(0); - setupParams->algoConfigs.resize(0); - for (SizeType32 bi = 0; bi < requestSize; bi++) - { - SizeType32 gbi = requestIds[bi]; - setupParams->prompt.emplace_back(mPrompt[gbi]); - setupParams->algoConfigs.emplace_back(mTestParam.w, mTestParam.n, mTestParam.g); - PRINT_TOKENS(setupParams->prompt[bi]); - setupParams->generationLengths = mGenerationLengths; - setupParams->positionOffsets = mPositionOffsets; - setupParams->attentionPackedMasks = mPackedMasks; - } - std::vector seed(requestIds.begin(), requestIds.end()); - setupParams->randomSeed = std::make_optional(seed); - TensorPtr newRequestSlots = ITensor::slice(mBatchSlotsMax, batchSize, requestSize); - PRINT_VALUES(newRequestSlots); - PRINT_VALUES(mBatchSlotsMax); - mBatchSlots = ITensor::slice(mBatchSlotsMax, 0, batchSize); - mDecodingWorkspace->setDeviceBatchSlots(newRequestSlots); - mDecoder->setup(requestSize, beamSize, newRequestSlots, setupParams, mDecodingWorkspace); - - PRINT_VALUES(mPositionOffsets); - - batchSize += requestIds.size(); - mBatchSlots = ITensor::slice(mBatchSlotsMax, 0, batchSize); - TLLM_LOG_DEBUG("new Requests mBatchSlots %s", D(mBatchSlots).values().c_str()); - PRINT_VALUES(mSequenceLengths); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadDecodingLayerTest::manageBatch() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const maxBatchSize = mTestParam.maxBatchSize; - auto requests = mBatchSlotsManager->alloc(); - if (requests.size() > 0) - { - newRequests(requests); - } - PRINT_VALUES(mSequenceLengths); - - auto batchSize = ITensor::volume(mBatchSlots->getShape()); - BufferRange batchSlotsRange(*mBatchSlots); - auto batchShape1D = ITensor::makeShape({batchSize}); - auto batchShape2D = ITensor::makeShape({batchSize, mMaxTokensPerStep}); - auto newBatchSize = 0; - PRINT_VALUES(mBatchSlots); - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - SizeType32 gbi = batchSlotsRange[bi]; - SizeType32 nbi = newBatchSize; - - TensorPtr theSequence = ITensor::at(mOutputIds, {gbi, 0}); - BufferRange theSequenceRange(*theSequence); - auto theSequenceLength = BufferRange(*mSequenceLengths)[gbi]; - auto theNumNewTokens = BufferRange(*mNumNewTokens)[gbi]; - - TensorPtr generated = ITensor::slice(theSequence, 0, theSequenceLength); - - PRINT_TOKENS(generated); - EXPECT_TRUE(mLlm[gbi]->verify(0, generated)); - - BufferRange(*mHistogram[gbi])[theNumNewTokens] += 1; - - if (BufferLocation(*theSequence).at(theSequenceLength - 1) == mAscii->getEndToken()) - { - TLLM_LOG_DEBUG("request[%d] ends: '%s'", gbi, D(theSequence).string().c_str()); - mScoreBoard[gbi] << "[" << gbi << "] ends. " << D(mHistogram[gbi]).values(); - mReports.push_back(mScoreBoard[gbi].str()); - mScoreBoard[gbi].str(""); - mScoreBoard[gbi].clear(); - mBatchSlotsManager->free(gbi); - } - else - { - batchSlotsRange[newBatchSize++] = gbi; - } - - auto theDraftLen = BufferRange(*mDraftLengths)[gbi]; - auto theGenerationLength = BufferRange(*mGenerationLengths)[gbi]; - TLLM_CHECK_DEBUG_WITH_INFO( - theDraftLen + 1 == theGenerationLength, "%d + 1 == %d", theDraftLen, theGenerationLength); - BufferLocation(*mTokensPerStep).at(gbi) = theGenerationLength; - - BufferLocation(*mInputTokensBatch).at(nbi, 0) = theSequenceRange[theSequenceLength - 1]; - mBufferManager->copy(*ITensor::slice(mDraftTokens, {gbi, 0}, theDraftLen), - *ITensor::slice(mInputTokensBatch, {nbi, 1}, theDraftLen)); - mBufferManager->copy(*ITensor::slice(mPositionIds, {gbi, 0}), *ITensor::slice(mPositionIdsBatch, {nbi, 0})); - BufferLocation(*mPositionIdsBatch).at(nbi, 0) = theSequenceLength - 1; - - TLLM_LOG_DEBUG("W=%d, N=%d, G=%d, w=%d, n=%d, g=%d, draftLen = %d", mTestParam.maxW, mTestParam.maxN, - mTestParam.maxG, mTestParam.w, mTestParam.n, mTestParam.g, theDraftLen); - - auto len = BufferRange(*mTokensPerStep)[gbi]; - PRINT_TOKENS(ITensor::slice(mInputTokensBatch, {nbi, 0}, len)); - PRINT_VALUES(ITensor::slice(mPositionIdsBatch, {nbi, 0}, len)); - } - mBatchSlots = ITensor::slice(mBatchSlotsMax, 0, newBatchSize); - PRINT_VALUES(mBatchSlots); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void convertInt32ToBool(TensorPtr const& dst, TensorConstPtr const& src) -{ - auto dstShape = dst->getShape(); - auto srcShape = src->getShape(); - TLLM_CHECK(dstShape.d[0] == srcShape.d[0]); - TLLM_CHECK(dstShape.d[1] <= srcShape.d[1] * 32); - BufferLocation dstLocation(*dst); - BufferLocation srcLocation(*src); - auto testBit = [](SizeType32 x, SizeType32 idx) { return x & (1 << idx); }; - for (auto i = 0; i < dstShape.d[0]; i++) - { - for (auto j = 0; j < dstShape.d[1]; j++) - { - dstLocation.at(i, j) = testBit(srcLocation.at(i, j / 32), j % 32); - } - } -} - -void convertBoolToInt32(TensorPtr const& dst, TensorConstPtr const& src) -{ - auto dstShape = dst->getShape(); - auto srcShape = src->getShape(); - TLLM_CHECK(dstShape.d[0] == srcShape.d[0]); - TLLM_CHECK(dstShape.d[1] * 32 >= srcShape.d[1]); - BufferLocation dstLocation(*dst); - BufferLocation srcLocation(*src); - - for (auto i = 0; i < dstLocation.size(); i++) - { - dstLocation[i] = 0; - } - - auto setBit = [](SizeType32& x, SizeType32 idx, bool value) { x |= (value << idx); }; - for (auto i = 0; i < srcShape.d[0]; i++) - { - for (auto j = 0; j < srcShape.d[1]; j++) - { - setBit(dstLocation.at(i, j / 32), j % 32, srcLocation.at(i, j)); - } - } -} - -void LookaheadDecodingLayerTest::llmForward() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto batchSize = ITensor::volume(mBatchSlots->getShape()); - - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - auto gbi = BufferRange(*mBatchSlots)[bi]; - auto start = BufferRange(*mSequenceLengths)[gbi] - 1; - auto len = BufferRange(*mTokensPerStep)[gbi]; - TLLM_LOG_DEBUG("LookaheadDecodingLayerTest::llmForward input len=%d", len); - TensorPtr output = ITensor::slice(mProbs, {bi, 0}, len); - TensorPtr golden = ITensor::slice(mGoldenSampledTokens, {gbi, 0}, len); - - BufferRange idRange(*ITensor::slice(mPositionIdsBatch, {bi, 0}, len)); - BufferRange offsetRange(*ITensor::slice(mPositionOffsets, {gbi, 0}, len)); - PRINT_VALUES(ITensor::slice(mPositionIdsBatch, {bi, 0})); - PRINT_VALUES(ITensor::slice(mPositionOffsets, {bi, 0})); - for (auto i = 0; i < idRange.size(); i++) - { - TLLM_CHECK(idRange[i] == start + offsetRange[i]); - } - - if (false) - { - convertInt32ToBool(ITensor::at(mPackedMasksBool, {gbi}), ITensor::at(mPackedMasks, {gbi})); - mLlm[gbi]->forward(output, // - ITensor::slice(mInputTokensBatch, {bi, 0}, len), // - ITensor::slice(mPositionIdsBatch, {bi, 0}, len), // - ITensor::at(mPackedMasksBool, {gbi})); - } - else - { - convertInt32ToBool(ITensor::at(mPackedMasksBool, {gbi}), ITensor::at(mPackedMasks, {gbi})); - mLlm[gbi]->forward(output, // - start, // - ITensor::slice(mInputTokensBatch, {bi, 0}, len), // - ITensor::slice(mPositionOffsets, {gbi, 0}, len), // - ITensor::at(mPackedMasksBool, {gbi})); - } - - mAscii->logitsToTensor(golden, output); - TLLM_LOG_DEBUG("batch[%d] LLM golden: '%s'", gbi, D(golden).tokens().c_str()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadDecodingLayerTest::decodeForward() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto batchSize = ITensor::volume(mBatchSlots->getShape()); - PRINT_VALUES(mBatchSlots); - - auto inputParams = std::make_shared(mEndIds, mBatchSlots); - inputParams->localBatchSize = batchSize; - inputParams->logits = ITensor::slice(mProbs, 0, batchSize); - inputParams->batchSlots = mBatchSlots; - inputParams->curTokensPerStep = mTokensPerStep; - - auto outputParams = std::make_shared(mOutputIds); - - PRINT_VALUES(mSequenceLengths); - outputParams->sequenceLength = mSequenceLengths; - outputParams->nextDraftLengths = mDraftLengths; - outputParams->prevDraftLengths = mPrevDraftLengths; - outputParams->nextDraftTokens = mDraftTokens; - outputParams->packedMasks = mPackedMasks; - outputParams->numNewTokens = mNumNewTokens; - outputParams->newTokens = mNewTokens; - outputParams->numNewTokensCumSum = mNumNewTokensCumSum; - outputParams->pathsOffsets = mPathsOffsets; - outputParams->generationLengths = mGenerationLengths; - outputParams->positionOffsets = mPositionOffsets; - outputParams->positionIds = mPositionIds; - outputParams->packedMasks = mPackedMasks; - - PRINT_VALUES(mTokensPerStep); - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mDecoder->forwardAsync(outputParams, inputParams, mDecodingWorkspace); - - mStream->synchronize(); - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mDecoder->forwardSync(outputParams, inputParams, mDecodingWorkspace); - - mStream->synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadDecodingLayerTest::verifyDecode() -{ - auto batchSize = ITensor::volume(mBatchSlots->getShape()); - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - auto gbi = BufferRange(*mBatchSlots)[bi]; - auto len = BufferRange(*mTokensPerStep)[gbi]; - auto sequenceLength = BufferLocation(*mSequenceLengths).at(gbi); - - auto draftLength = BufferLocation(*mDraftLengths).at(gbi); - auto generationLength = BufferLocation(*mGenerationLengths).at(gbi); - BufferRange posOffsetRange(*ITensor::slice(mPositionOffsets, {gbi, 0}, generationLength)); - BufferRange posIdRange(*ITensor::slice(mPositionIds, {gbi, 0}, generationLength)); - TLLM_LOG_DEBUG("generationLength = %d, draftLength = %d", generationLength, draftLength); - TLLM_CHECK(draftLength + 1 == generationLength); - TLLM_CHECK(posOffsetRange[0] == 0); - TLLM_CHECK(posIdRange[0] == sequenceLength - 1); - for (SizeType32 i = 0; i < posIdRange.size(); i++) - { - TLLM_CHECK(posIdRange[i] == posOffsetRange[i] + sequenceLength - 1); - } - } - - BufferRange cumSumRange(*mNumNewTokensCumSum); - BufferRange pathOffsetsRange(*mPathsOffsets); - PRINT_VALUES(mNumNewTokensCumSum); - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - auto gbi = BufferRange(*mBatchSlots)[bi]; - SizeType32 pathOffsetBegin = cumSumRange[bi]; - SizeType32 pathOffsetEnd = cumSumRange[bi + 1]; - TensorPtr golden = ITensor::at(mGoldenSampledTokens, {gbi}); - auto sequenceLength = BufferLocation(*mSequenceLengths).at(gbi); - auto numNewTokens = BufferLocation(*mNumNewTokens).at(gbi); - TensorPtr newTokens = ITensor::slice(mOutputIds, {gbi, 0, sequenceLength - numNewTokens}, numNewTokens); - BufferRange goldenRange(*ITensor::at(mGoldenSampledTokens, {gbi})); - BufferRange newTokensRange(*newTokens); - - SizeType32 ni = 1; - for (SizeType32 poi = pathOffsetBegin; poi < pathOffsetEnd; poi++) - { - TLLM_CHECK(goldenRange[pathOffsetsRange[poi] + 1] == newTokensRange[ni++]); - } - } -} - -void LookaheadDecodingLayerTest::runTest(TestParam const& param) -{ - TLLM_LOG_DEBUG("TEST BEGIN: maxBatchSize=%d, mode=%d, WNG=(%d, %d, %d), wng=(%d, %d, %d)", param.maxBatchSize, - param.batchType, param.maxW, param.maxN, param.maxG, param.w, param.n, param.g); - srand(42); - - mTestParam = param; - allocateBuffers(); - - int step = 0; - for (; !mBatchSlotsManager->finished() && step < 3000; step++) - { - TLLM_LOG_DEBUG("!!!!!!!!!!!!!!!! < %d > !!!!!!!!!!!!!!!!", step); - manageBatch(); - if (ITensor::volume(mBatchSlots->getShape())) - { - llmForward(); - mStream->synchronize(); - decodeForward(); - verifyDecode(); - } - } - - for (auto& r : mReports) - { - TLLM_LOG_DEBUG(r); - } - if (!mBatchSlotsManager->finished()) - { - TLLM_LOG_INFO("step=%d is not enough", step); - } -} - -TEST_F(LookaheadDecodingLayerTest, singleOnce) -{ - this->runTest(TestParam{16, TestParam::SINGLE_ONCE, 5, 3, 5, 3, 5, 3}); -} - -TEST_F(LookaheadDecodingLayerTest, singleTwice) -{ - this->runTest(TestParam{16, TestParam::SINGLE_TWICE, 7, 5, 7, 5, 7, 5}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamic) -{ - this->runTest(TestParam{16, TestParam::DYNAMIC, 5, 5, 5, 5, 5, 5}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamicLarge) -{ - this->runTest(TestParam{32, TestParam::DYNAMIC, 7, 6, 7, 6, 9, 8}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamicSmall_110) -{ - this->runTest(TestParam{16, TestParam::SINGLE_TWICE, 1, 1, 2, 2, 0, 0}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamicSmall_311) -{ - this->runTest(TestParam{32, TestParam::DYNAMIC, 3, 2, 2, 2, 1, 1}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamicSmall_131) -{ - this->runTest(TestParam{32, TestParam::DYNAMIC, 1, 1, 3, 2, 1, 1}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamicSmall_113) -{ - this->runTest(TestParam{32, TestParam::DYNAMIC, 1, 1, 2, 2, 3, 2}); -} - -TEST_F(LookaheadDecodingLayerTest, dynamicSmall_112110) -{ - this->runTest(TestParam{4, TestParam::SINGLE_TWICE, 1, 1, 2, 1, 1, 0}); -} - -using ParamType = std::tuple, - std::tuple, std::tuple>; - -static int g_id = 0; - -std::string generateTestName(testing::TestParamInfo const& info) -{ - auto [maxBatchSize, mode, Ww, Nn, Gg] = info.param; - auto [W, w] = Ww; - auto [N, n] = Nn; - auto [G, g] = Gg; - std::ostringstream buf; - buf << (g_id++) << "maxBatchSize_" << maxBatchSize << "__mode_" << mode << '_' << '_' << W << '_' << w << '_' << '_' - << N << '_' << n << '_' << '_' << G << '_' << g << '_'; - return buf.str(); -} - -class ParamTest : public LookaheadDecodingLayerTest, public ::testing::WithParamInterface -{ -}; - -TEST_P(ParamTest, Test) -{ - srand(42); - - auto [maxBatchSize, mode, Ww, Nn, Gg] = GetParam(); - auto [W, w] = Ww; - auto [N, n] = Nn; - auto [G, g] = Gg; - if (!executor::LookaheadDecodingConfig::isLegal(W, N, G) || !executor::LookaheadDecodingConfig::isLegal(w, n, g)) - { - TLLM_LOG_DEBUG("Just Pass for illegal parameter combination"); - GTEST_SKIP() << "Algorithm does not support these parameters WNG=(" << W << ", " << N << ", " << G << "), wng=(" - << w << ", " << n << ", " << g << ")"; - } - runTest(TestParam{maxBatchSize, mode, W, w, N, n, G, g}); -} - -INSTANTIATE_TEST_SUITE_P(LookaheadDecodingLayerParamTest, ParamTest, - testing::Combine( // - testing::Values(4, 16), testing::Values(TestParam::DYNAMIC), - testing::Values(std::make_tuple(1, 1), std::make_tuple(3, 3), std::make_tuple(5, 5), std::make_tuple(2, 1), - std::make_tuple(3, 2), std::make_tuple(5, 3)), - testing::Values(std::make_tuple(1, 1), std::make_tuple(3, 3), std::make_tuple(5, 5), std::make_tuple(2, 1), - std::make_tuple(3, 2), std::make_tuple(5, 3)), - testing::Values(std::make_tuple(0, 0), std::make_tuple(3, 3), std::make_tuple(5, 5), std::make_tuple(1, 0), - std::make_tuple(3, 2), std::make_tuple(5, 3))), - generateTestName); - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/lookaheadPoolManagerTest.cpp b/cpp/tests/unit_tests/layers/lookaheadPoolManagerTest.cpp deleted file mode 100644 index 14e52b808359..000000000000 --- a/cpp/tests/unit_tests/layers/lookaheadPoolManagerTest.cpp +++ /dev/null @@ -1,174 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include - -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/layers/lookaheadPoolManager.h" -#include "tests/unit_tests/layers/randomLlm.h" - -namespace tensorrt_llm::tests::layers -{ -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using TensorPtr = runtime::ITensor::SharedPtr; -using TensorConstPtr = runtime::ITensor::SharedConstPtr; - -void printMap( - char const* name, std::unordered_map> const& tokenMap) -{ - std::ostringstream buf; - buf << name << std::endl; - for (auto const& [key, value] : tokenMap) - { - buf << static_cast(key) << ": "; - for (auto const& tup : value) - { - buf << "("; - for (auto const& token : BufferRange(*tup)) - { - buf << static_cast(token) << ","; - } - buf << "),"; - } - buf << std::endl; - } - TLLM_LOG_DEBUG(buf.str()); -} - -bool isTensorEqString(TensorConstPtr const& a, std::string b) -{ - TLLM_CHECK(ITensor::volume(a->getShape()) == static_cast(b.size())); - auto ar = BufferRange(*a); - return std::equal(ar.begin(), ar.end(), b.begin()); -} - -TEST(LookaheadPoolManagerTest, fillAndUpdate) -{ - SizeType32 constexpr W{5}; - SizeType32 constexpr N{4}; - SizeType32 constexpr G{5}; - auto prompt = initTensor("hello world; hello world. live is life."); - LookaheadPoolManager pm(G); - pm.setup(G); - pm.accept(prompt, N); - printMap("Token map after fill with prompt", pm.getMap()); - /*** - s: ( ,l,i,), - v: (e, ,i,), - i: (v,e, ,),(s, ,l,),(f,e,.,), - d: (;, ,h,),(., ,l,), - w: (o,r,l,), - : (h,e,l,),(w,o,r,),(l,i,v,),(i,s, ,),(l,i,f,), - .: ( ,l,i,), - ;: ( ,h,e,), - o: ( ,w,o,),(r,l,d,), - l: (l,o, ,),(o, ,w,),(d,., ,),(i,v,e,),(i,f,e,), - r: (l,d,;,),(l,d,.,), - e: (l,l,o,),( ,i,s,), - h: (e,l,l,), - **/ - - LookaheadPoolManager::Key lastToken = 'l'; - auto list = pm.guess(lastToken, G); - for (auto const& ngram : list) - { - PRINT_TOKENS(ngram); - } - /*** - l: (l,o, ,),(o, ,w,),(d,., ,),(i,v,e,),(i,f,e,), - **/ - - auto pastTokens = initTensor(std::string("abcde12345hijkm"), ITensor::makeShape({5, 3})); - auto keyTokens = initTensor(std::string("lvwxy")); - pm.update(keyTokens, pastTokens); - printMap("Token map after update", pm.getMap()); - /** Noted, we update the map with N=4, so the map has different sizes of ngrams. - y: (j,k,m,), - x: (5,h,i,), - e: (l,l,o,),( ,i,s,), - r: (l,d,;,),(l,d,.,), - l: (o, ,w,),(d,., ,),(i,v,e,),(i,f,e,),(a,b,c,), - o: ( ,w,o,),(r,l,d,), - ;: ( ,h,e,), - h: (e,l,l,), - .: ( ,l,i,), - : (h,e,l,),(w,o,r,),(l,i,v,),(i,s, ,),(l,i,f,), - w: (o,r,l,),(2,3,4,), - d: (;, ,h,),(., ,l,), - i: (v,e, ,),(s, ,l,),(f,e,.,), - v: (e, ,i,),(d,e,1,), - s: ( ,l,i,), - */ - - lastToken = 'w'; - list = pm.guess(lastToken, G); - for (auto const& ngram : list) - { - PRINT_TOKENS(ngram); - } - /** - w: (o,r,l,),(2,3,4,), - */ - - ASSERT_EQ(list.size(), 2); - auto it = list.begin(); - EXPECT_TRUE(isTensorEqString(*it, "orl")); - it++; - EXPECT_TRUE(isTensorEqString(*it, "234")); - - pastTokens = initTensor(std::string("dogde12345hijkm"), ITensor::makeShape({5, 3})); - pm.update(keyTokens, pastTokens); - - pastTokens = initTensor(std::string("catde12345hijkm"), ITensor::makeShape({5, 3})); - pm.update(keyTokens, pastTokens); - - pastTokens = initTensor(std::string("abcde12345hijkm"), ITensor::makeShape({5, 3})); - pm.update(keyTokens, pastTokens); - - printMap("Token map after update more for key 'l'", pm.getMap()); - /** - y: (j,k,m,), - x: (5,h,i,), - e: (l,l,o,),( ,i,s,), - r: (l,d,;,),(l,d,.,), - l: (i,v,e,),(i,f,e,),(d,o,g,),(c,a,t,),(a,b,c,), - o: ( ,w,o,),(r,l,d,), - ;: ( ,h,e,), - h: (e,l,l,), - .: ( ,l,i,), - : (h,e,l,),(w,o,r,),(l,i,v,),(i,s, ,),(l,i,f,), - w: (o,r,l,),(2,3,4,), - d: (;, ,h,),(., ,l,), - i: (v,e, ,),(s, ,l,),(f,e,.,), - v: (e, ,i,),(d,e,1,), - s: ( ,l,i,), - */ - lastToken = 'l'; - list = pm.guess(lastToken, G); - ASSERT_EQ(list.size(), G); - it = list.begin(); - EXPECT_TRUE(isTensorEqString(*it, "ive")); - it++; - EXPECT_TRUE(isTensorEqString(*it, "ife")); - it++; - EXPECT_TRUE(isTensorEqString(*it, "dog")); - it++; - EXPECT_TRUE(isTensorEqString(*it, "cat")); - it++; - EXPECT_TRUE(isTensorEqString(*it, "abc")); -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp b/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp deleted file mode 100644 index 4ca7206c3bc8..000000000000 --- a/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp +++ /dev/null @@ -1,242 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/layers/lookaheadAlgorithm.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tests/unit_tests/layers/randomLlm.h" - -namespace tensorrt_llm::tests::layers -{ -namespace tk = tensorrt_llm::kernels; -namespace trk = tensorrt_llm::runtime::kernels; -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; - -using TensorPtr = runtime::ITensor::SharedPtr; - -TEST(LookaheadRandomllm, forward) -{ - auto ascii = std::make_shared(); - EXPECT_EQ(ascii->getVocabSize(), 128); - { - auto tensor = ascii->tokenToLogits(static_cast('a')); - auto token = ascii->logitsToToken(tensor); - EXPECT_EQ(static_cast(token), 'a'); - } - { - auto tensor = ascii->tokenToLogits(static_cast('W')); - auto token = ascii->logitsToToken(tensor); - EXPECT_EQ(static_cast(token), 'W'); - } - { - std::string str("hello world!"); - TensorPtr logits - = BufferManager::cpu(ITensor::makeShape({static_cast(str.size()), ascii->getVocabSize()}), - tensorrt_llm::DataType::kFLOAT); - ascii->stringToLogits(logits, str); - auto result = ascii->logitsToString(logits); - EXPECT_EQ(result, str); - } - - std::string oracle( - "The following example uses a lambda-expression to increment all of the elements of a vector and " - "then uses an overloaded operator() in a function object (a.k.a., \"functor\") to compute their sum. Note that " - "to compute the sum, it is recommended to use the dedicated algorithm std::accumulate."); - LookaheadRandomLlm llm(ascii, oracle); - { - TLLM_LOG_DEBUG("oracle[22]='%c'", oracle[22]); - std::string input("ubcs23eess a la"); - auto len = static_cast(input.size()); - TensorPtr inputTokens = initTensor(input); - std::vector positionIdVec({22, 23, 24, 23, 24, 25, 24, 25, 26, 25, 26, 27, 26, 27, 28}); - TensorPtr positionIds = ITensor::wrap(positionIdVec, ITensor::makeShape({len})); - TensorPtr outputLogits - = BufferManager::cpu(ITensor::makeShape({len, ascii->getVocabSize()}), tensorrt_llm::DataType::kFLOAT); - - llm.forward(outputLogits, inputTokens, positionIds); - - auto result = ascii->logitsToString(outputLogits); - auto invalid = ascii->getInvalidToken(); - TLLM_LOG_DEBUG("result=%s", result.c_str()); - for (SizeType32 i = 0; i < len; i++) - { - if (result[i] != invalid) - { - EXPECT_EQ(result[i], oracle[positionIdVec[i] + 1]); - } - } - } -} - -TEST(LookaheadRandomllm, gpuSampling) -{ - auto mStream = std::make_shared(); - auto mBufferManager = std::make_shared(mStream); - - int32_t device; - struct cudaDeviceProp mDeviceProp; - cudaGetDevice(&device); - cudaGetDeviceProperties(&mDeviceProp, device); - - // auto mAscii = std::make_shared(); - auto mAscii = std::make_shared(); - - std::vector text({std::string("0123456789abcdef0123456789abcdef0123456&"), - std::string("hello world, hello world, hello world!!&"), - std::string("To be or not to be that is the question&"), - std::string("To be or not to be that is the question&")}); - - SizeType32 W = 5, N = 5, G = 5; - SizeType32 maxBatchSize = 16; - std::vector batchSlotsVec({1, 4, 7, 11}); - SizeType32 batchSize = batchSlotsVec.size(); - SizeType32 vocabSizePadded = mAscii->getVocabSize(); - SizeType32 vocabSize = vocabSizePadded; - SizeType32 maxTokensPerStep = (W + G) * (N - 1); - SizeType32 maxNumHeads = 1; - SizeType32 mRuntimeMaxTopK = 1; - SizeType32 mMaxTopK = 1; - SizeType32 mMaxTopP = 1.0; - - auto maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - auto maxBatchShape3D = ITensor::makeShape({maxBatchSize, maxTokensPerStep, vocabSize}); - auto batchShape1D = ITensor::makeShape({batchSize}); - - uint32_t mSeed = 0; - SizeType32 mMaxSeqLen = 128; - - SizeType32 workspaceSize - = tensorrt_llm::kernels::getTopKWorkspaceSize(maxBatchSize, maxTokensPerStep, mMaxTopK, vocabSizePadded); - TensorPtr workspaceDevice = mBufferManager->pinned( - ITensor::makeShape({static_cast(workspaceSize)}), tensorrt_llm::DataType::kINT8); - - auto const dataType = TRTDataType::value; - auto const ptrType = TRTDataType::value; - - // Allocate GPU data - TensorPtr mSeqLengths = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - TensorPtr mFinished = BufferManager::pinned(maxBatchShape1D, TRTDataType::value); - TensorPtr mEndIds = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - TensorPtr mTopPs = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kFLOAT); - TensorPtr mTopKs = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - TensorPtr mSkipDecode = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kBOOL); - TensorPtr mTokensPerStep = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - - TensorPtr mCurandStates = BufferManager::pinned( - ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); - TensorPtr mOutputIds - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - - TensorPtr mProbs = BufferManager::pinned(maxBatchShape3D, dataType); - - TensorPtr mBatchSlots = BufferManager::pinned(batchShape1D, tensorrt_llm::DataType::kINT32); - - ///////////////////////////////////// - std::copy(batchSlotsVec.begin(), batchSlotsVec.end(), BufferRange(*mBatchSlots).begin()); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - // Allocate and init curand states - tk::invokeCurandInitialize(reinterpret_cast(bufferCast(*mCurandStates)), batchSlotsPtr, - batchSize, mSeed, mStream->get()); - - // Init by zero. - trk::invokeFill(*mFinished, uint8_t{0}, *mStream); - trk::invokeFill(*mOutputIds, int32_t{0}, *mStream); - trk::invokeFill(*mSkipDecode, false, *mStream); - trk::invokeFill(*mEndIds, mAscii->getEndToken(), *mStream); - trk::invokeFill(*mTopPs, float{1.0}, *mStream); - trk::invokeFill(*mTopKs, int32_t{1}, *mStream); - trk::invokeFill(*mSeqLengths, int32_t{0}, *mStream); - trk::invokeFill(*mTokensPerStep, maxTokensPerStep, *mStream); - - TLLM_CHECK(mMaxTopK * maxTokensPerStep <= mMaxSeqLen); - - // Init logits randomly - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - TensorPtr one = ITensor::at(mProbs, {bi}); - mAscii->stringToLogits(one, text[bi]); - auto result = mAscii->logitsToString(one); - EXPECT_EQ(result, text[bi]); - } - - tensorrt_llm::kernels::TopKSamplingKernelParams kernelParams; - kernelParams.logProbs = bufferCast(*mProbs); - kernelParams.logProbsPtrs = nullptr; - // kernelParams.outputIdsPtrs = bufferCast(*mIdsPtrHost); - // kernelParams.outputIds = nullptr; - kernelParams.outputIdsPtrs = nullptr; - kernelParams.outputIds = bufferCast(*mOutputIds); - kernelParams.maxSeqLen = mMaxSeqLen; - kernelParams.workspace = workspaceDevice->data(); - kernelParams.maxTopP = 1.0; - kernelParams.topPs = bufferCast(*mTopPs); - kernelParams.maxTopK = mMaxTopK; - kernelParams.topKs = bufferCast(*mTopKs); - kernelParams.sequenceLengths = bufferCast(*mSeqLengths); - kernelParams.endIds = bufferCast(*mEndIds); - kernelParams.batchSlots = bufferCast(*mBatchSlots); - kernelParams.finishedInput = reinterpret_cast( - bufferCast(*mFinished)); - kernelParams.finishedOutput = reinterpret_cast( - bufferCast(*mFinished)); - kernelParams.skipDecode = bufferCast(*mSkipDecode); - kernelParams.cumLogProbs = nullptr; - kernelParams.outputLogProbs = nullptr; - kernelParams.curandState = reinterpret_cast(bufferCast(*mCurandStates)); - kernelParams.batchSize = batchSize; - kernelParams.maxBatchSize = maxBatchSize; - kernelParams.maxTokensPerStep = maxTokensPerStep; - kernelParams.tokensPerStep = bufferCast(*mTokensPerStep); - kernelParams.vocabSizePadded = vocabSize; - kernelParams.normalizeLogProbs = false; - kernelParams.logitsHasProbs = false; - kernelParams.returnAllSelectedTokens = false; - - PRINT_TOKENS(mEndIds); - PRINT_VALUES(mTokensPerStep); - PRINT_VALUES(mBatchSlots); - PRINT_VALUES(mTopKs); - tensorrt_llm::kernels::invokeBatchTopKSampling(kernelParams, mStream->get()); - - mStream->synchronize(); - - std::ostringstream buf; - buf << "finished states: "; - for (SizeType32 bi = 0; bi < maxBatchSize; bi++) - { - buf << "[" << bi << "]=" << kernelParams.finishedOutput[bi].isFinished() << ", "; - } - TLLM_LOG_DEBUG(buf.str()); - - for (SizeType32 bi = 0; bi < batchSize; bi++) - { - SizeType32 gbi = kernelParams.batchSlots[bi]; - bool finished = kernelParams.finishedOutput[bi].isFinished(); - TensorPtr one = ITensor::at(mOutputIds, {gbi}); - auto oneRange = BufferRange(*one); - std::vector result(mMaxSeqLen, '\0'); - std::copy(oneRange.begin(), oneRange.end(), result.begin()); - TLLM_LOG_DEBUG(result.data()); - EXPECT_EQ(text[bi], result.data()); - } -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp b/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp deleted file mode 100644 index a93955d02c9a..000000000000 --- a/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp +++ /dev/null @@ -1,522 +0,0 @@ -/* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "medusaDecodeLayerTest.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/runtime/medusaModule.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include - -namespace tensorrt_llm::tests::layers -{ - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::layers; -using namespace tensorrt_llm::common; - -namespace tk = tensorrt_llm::kernels; -namespace trk = tensorrt_llm::runtime::kernels; - -constexpr float EPSILON = 1e-20f; - -template -void MedusaDecodingLayerTest::SetUp() -{ - mStream = std::make_shared(); - mBufferManager = std::make_shared(mStream); -} - -template -void MedusaDecodingLayerTest::allocateBuffers() -{ - auto speculativeDecodingModule = std::make_shared(mMaxDraftPathLen, mMaxDecodingTokens - 1); - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain( - mMaxBatchSize, 1, mVocabSize, mVocabSizePadded, speculativeDecodingModule); - mMedusaDecodingLayer - = std::make_shared>(decodingDomain, mBufferManager); - - auto const dataType = TRTDataType::value; - - // clang-format off - - // prob = (0.0, 0.0, 0.0, 0.0, 0.4, 0.3, 0.2, 0.1, 0.0) - std::vector targetLogitsInit = { - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 0 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 1 - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 2 - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 3 - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 4 - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 5 - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 6 - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 7 - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 8 - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 9 - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 10 - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX // token 11 - }; - // Sampled tokens with K=1 - // [4, 0, 2, 1, 3, 4, 3, 0, 2, 3, 4, 1] - - std::vector medusaLogitsInit = { - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 0 head=0 ids: [4, 5, 6, 7] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 1 head=0 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 2 head=0 ids: [2, 3, 4, 5] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 3 head=0 ids: [1, 2, 3, 4] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 4 head=0 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 5 head=0 ids: [4, 5, 6, 7] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 6 head=0 ids: [3, 4, 5, 6] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 7 head=0 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 8 head=0 ids: [2, 3, 4, 5] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 9 head=0 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 10 head=0 ids: [4, 5, 6, 7] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 11 head=0 ids: [1, 2, 3, 4] - - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 0 head=1 ids: [2, 3, 4, 5] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 1 head=1 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 2 head=1 ids: [4, 5, 6, 7] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 3 head=1 ids: [1, 2, 3, 4] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 4 head=1 ids: [4, 5, 6, 7] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 5 head=1 ids: [2, 3, 4, 5] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 6 head=1 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 7 head=1 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 8 head=1 ids: [3, 4, 5, 6] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 9 head=1 ids: [0, 1, 2, 3] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 10 head=1 ids: [1, 2, 3, 4] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 11 head=1 ids: [4, 5, 6, 7] - - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 0 head=2 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 1 head=2 ids: [4, 5, 6, 7] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 2 head=2 ids: [1, 2, 3, 4] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 3 head=2 ids: [2, 3, 4, 5] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 4 head=2 ids: [4, 5, 6, 7] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 5 head=2 ids: [3, 4, 5, 6] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 6 head=2 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 7 head=2 ids: [2, 3, 4, 5] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 8 head=2 ids: [1, 2, 3, 4] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 9 head=2 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 10 head=2 ids: [4, 5, 6, 7] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 11 head=2 ids: [3, 4, 5, 6] - - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 0 head=3 ids: [4, 5, 6, 7] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 1 head=3 ids: [4, 5, 6, 7] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 2 head=3 ids: [1, 2, 3, 4] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 3 head=3 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, // token 4 head=3 ids: [4, 5, 6, 7] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 5 head=3 ids: [2, 3, 4, 5] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 6 head=3 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 7 head=3 ids: [3, 4, 5, 6] - -FLT_MAX, -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, // token 8 head=3 ids: [3, 4, 5, 6] - -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 9 head=3 ids: [1, 2, 3, 4] - -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, // token 10 head=3 ids: [0, 1, 2, 3] - -FLT_MAX, -FLT_MAX, -0.9163, -1.2040, -1.6094, -2.3026, -FLT_MAX, -FLT_MAX, -FLT_MAX // token 11 head=3 ids: [2, 3, 4, 5] - }; - - // clang-format on - - auto const targetLogitsHost - = ITensor::wrap(targetLogitsInit.data(), dataType, ITensor::makeShape({mMaxDecodingTokens, mVocabSizePadded})); - - TensorPtr medusaLogitsHost = ITensor::wrap(medusaLogitsInit.data(), dataType, - ITensor::makeShape({mMaxDraftPathLen, mMaxDecodingTokens, mVocabSizePadded})); - - mTargetLogitsDevice - = mBufferManager->gpu(ITensor::makeShape({mBatchSize, mMaxDecodingTokens, mVocabSizePadded}), dataType); - - mFinishedDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize}), TRTDataType::value); - - mOutputIdsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); - - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mPathsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens, mMaxDraftPathLen + 1}), tensorrt_llm::DataType::kINT32); - - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mAcceptedLengths = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mTreeIdsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens - 1}), tensorrt_llm::DataType::kINT32); - - mMedusaLogitsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxDraftPathLen, mMaxBatchSize, mMaxDecodingTokens, mVocabSizePadded}), dataType); - - mNextDraftTokensDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens - 1}), tensorrt_llm::DataType::kINT32); - - mTokensPerStepDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - - mAcceptedLengthCumSumDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize + 1}), tensorrt_llm::DataType::kINT32); - - mPackedPathsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize * mMaxDraftPathLen}), tensorrt_llm::DataType::kINT32); - - for (int32_t bi = 0; bi < mBatchSize; ++bi) - { - auto logitsDeviceView = ITensor::slice(mTargetLogitsDevice, bi, 1); - mBufferManager->copy(*targetLogitsHost, *logitsDeviceView); - } - - for (int32_t hi = 0; hi < mMaxDraftPathLen; ++hi) - { - TensorPtr logitsHeadDeviceView = ITensor::slice(mMedusaLogitsDevice, hi, 1); - TensorPtr logitsHeadHostView = ITensor::slice(medusaLogitsHost, hi, 1); - logitsHeadDeviceView->squeeze(0); - for (int32_t bi = 0; bi < mBatchSize; ++bi) - { - TensorPtr logitsHeadBatchDeviceView = ITensor::slice(logitsHeadDeviceView, bi, 1); - mBufferManager->copy(*logitsHeadHostView, *logitsHeadBatchDeviceView); - } - } - - mDecodingWorkspace = std::make_unique( - mBufferManager, decodingDomain, TRTDataType::value, mMedusaDecodingLayer->getWorkspaceSize()); -} - -template -void MedusaDecodingLayerTest::setup(SamplingParams& params) -{ - auto const endId = params.endId.value_or(mEndId); - trk::invokeFill(*mSeqLengthsDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mAcceptedLengths, SizeType32{0}, *mStream); - trk::invokeFill(*mFinishedDevice, uint8_t{0}, *mStream); - trk::invokeFill(*mOutputIdsDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mEndIdsDevice, TokenIdType{endId}, *mStream); - trk::invokeFill(*mNextDraftTokensDevice, TokenIdType{-1}, *mStream); - trk::invokeFill(*mPathsDevice, SizeType32{-1}, *mStream); - trk::invokeFill(*mTreeIdsDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mTokensPerStepDevice, SizeType32{0}, *mStream); - trk::invokeFill(*mAcceptedLengthCumSumDevice, SizeType32{-1}, *mStream); - trk::invokeFill(*mPackedPathsDevice, SizeType32{-1}, *mStream); - - auto batchSlotsPtr = bufferCast(*mBatchSlots); - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - batchSlotsPtr[bi] = 2 * bi; - } - - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const draftIdsHost = ITensor::wrap(reinterpret_cast(params.draftIds[bi].data()), - tensorrt_llm::DataType::kINT32, ITensor::makeShape({1, mMaxDecodingTokens - 1})); - auto draftIdsDeviceSlice = ITensor::slice(mNextDraftTokensDevice, batchSlotsPtr[bi], 1); - mBufferManager->copy(*draftIdsHost, *draftIdsDeviceSlice); - } - - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto& path = params.paths[bi]; - auto const numPaths = static_cast(params.paths[bi].size() / (mMaxDraftPathLen + 1)); - auto const pathsHost = ITensor::wrap(reinterpret_cast(path.data()), tensorrt_llm::DataType::kINT32, - ITensor::makeShape({1, numPaths, mMaxDraftPathLen + 1})); - TensorPtr pathsDeviceSlice = ITensor::slice(mPathsDevice, batchSlotsPtr[bi], 1); - pathsDeviceSlice->squeeze(0); - TensorPtr pathsNumPathsDeviceSlice = ITensor::slice(pathsDeviceSlice, 0, numPaths); - pathsNumPathsDeviceSlice->unsqueeze(0); - mBufferManager->copy(*pathsHost, *pathsNumPathsDeviceSlice); - } - - auto tokensPerStep = params.tokensPerStep; - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - TensorPtr tokensPerStepDeviceSlice = ITensor::slice(mTokensPerStepDevice, batchSlotsPtr[bi], 1); - trk::invokeFill(*tokensPerStepDeviceSlice, SizeType32{tokensPerStep[bi]}, *mStream); - } - - auto treeIds = params.treeIds; - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const tokensPerStep = static_cast(treeIds[bi].size()); - auto const treeIdsBatchHost = ITensor::wrap(treeIds[bi], ITensor::makeShape({tokensPerStep})); - TensorPtr treeIdsBatchDevice = ITensor::slice(mTreeIdsDevice, batchSlotsPtr[bi], 1); - treeIdsBatchDevice->squeeze(0); - auto const treeIdsBatchDeviceSlice = ITensor::slice(treeIdsBatchDevice, 0, tokensPerStep); - mBufferManager->copy(*treeIdsBatchHost, *treeIdsBatchDeviceSlice); - } - - auto setupParams = std::make_shared(); - setupParams->runtimeTopK = std::make_optional>(params.runtimeTopK); - setupParams->runtimeHeadsTopK = std::make_optional>>(params.runtimeHeadsTopK); - setupParams->randomSeed = {{0}}; - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mMedusaDecodingLayer->setup(mBatchSize, 1, mBatchSlots, setupParams, mDecodingWorkspace); - - mStream->synchronize(); -} - -template -std::shared_ptr MedusaDecodingLayerTest::createInputTensors() -{ - auto forwardParams = std::make_shared(mEndIdsDevice, mBatchSlots, mBatchSize); - - auto batchSlots = BufferRange(*mBatchSlots); - - forwardParams->logits = mTargetLogitsDevice; - - forwardParams->finished = mFinishedDevice; - - forwardParams->paths = mPathsDevice; - - forwardParams->treeIds = mTreeIdsDevice; - - std::vector> medusaLogits(mMaxBatchSize); - auto const medusaLogitsPtr = bufferCast(*mMedusaLogitsDevice); - for (SizeType32 bi = 0; bi < mMaxBatchSize; ++bi) - { - medusaLogits[bi].resize(mMaxDraftPathLen); - } - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - for (SizeType32 hi = 0; hi < mMaxDraftPathLen; ++hi) - { - TensorPtr logitsHead = ITensor::slice(mMedusaLogitsDevice, hi, 1); - logitsHead->squeeze(0); - TensorPtr logitsHeadBatch = ITensor::slice(logitsHead, bi, 1); - medusaLogits[batchSlots[bi]][hi] = logitsHeadBatch; - } - } - forwardParams->medusaLogits = medusaLogits; - - forwardParams->curTokensPerStep = mTokensPerStepDevice; - - forwardParams->targetTokensPerStep = mTokensPerStepDevice; - - return forwardParams; -} - -template -std::shared_ptr MedusaDecodingLayerTest::createOutputTensors() -{ - auto outputParams = std::make_shared(mOutputIdsDevice); - - outputParams->sequenceLength = mSeqLengthsDevice; - - outputParams->finished = mFinishedDevice; - - outputParams->nextDraftTokens = mNextDraftTokensDevice; - - outputParams->numNewTokens = mAcceptedLengths; - - outputParams->numNewTokensCumSum = mAcceptedLengthCumSumDevice; - - outputParams->pathsOffsets = mPackedPathsDevice; - - return outputParams; -} - -template -void MedusaDecodingLayerTest::checkResult(std::vector>> const& expectedOutTokens, - std::vector> const& expectedDraftTokens, std::vector const& finished, - SamplingParams& params) -{ - auto const nextDraftTokensHost = mBufferManager->copyFrom(*mNextDraftTokensDevice, runtime::MemoryType::kCPU); - auto const outputIdsHost = mBufferManager->copyFrom(*mOutputIdsDevice, runtime::MemoryType::kCPU); - auto const seqLenHost = mBufferManager->copyFrom(*mSeqLengthsDevice, runtime::MemoryType::kCPU); - auto const acceptedLengthsHost = mBufferManager->copyFrom(*mAcceptedLengths, runtime::MemoryType::kCPU); - auto const finishedHost = mBufferManager->copyFrom(*mFinishedDevice, runtime::MemoryType::kCPU); - auto const acceptedLengthCumSumHost - = mBufferManager->copyFrom(*mAcceptedLengthCumSumDevice, runtime::MemoryType::kCPU); - auto const packedPathsHost = mBufferManager->copyFrom(*mPackedPathsDevice, runtime::MemoryType::kCPU); - - mStream->synchronize(); - - auto nextDraftTokens = BufferRange(*nextDraftTokensHost); - auto outputIds = BufferRange(*outputIdsHost); - auto seqLen = BufferRange(*seqLenHost); - auto batchSlots = BufferRange(*mBatchSlots); - auto acceptedLengths = BufferRange(*acceptedLengthsHost); - auto acceptedLengthCumSum = BufferRange(*acceptedLengthCumSumHost); - auto packedPaths = BufferRange(*packedPathsHost); - auto finishedPtr - = reinterpret_cast(bufferCast(*finishedHost)); - - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto& expectedOutTokensBatch = expectedOutTokens[bi]; - auto const slot = batchSlots[bi]; - EXPECT_EQ(expectedOutTokensBatch.size(), seqLen[slot]); - EXPECT_EQ(expectedOutTokensBatch.size(), acceptedLengths[slot]); - for (SizeType32 ti = 0; ti < expectedOutTokensBatch.size(); ++ti) - { - EXPECT_GE(expectedOutTokensBatch[ti].count(outputIds[slot * mMaxSeqLen + ti]), 1); - } - EXPECT_EQ(acceptedLengthCumSum[bi], params.acceptedCumSum[bi]); - } - EXPECT_EQ(acceptedLengthCumSum[mBatchSize], params.acceptedCumSum[mBatchSize]); - for (SizeType32 ti = 0; ti < params.packedPaths.size(); ++ti) - { - EXPECT_EQ(packedPaths[ti], params.packedPaths[ti]); - } - - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto& expectedDraftTokensBatch = expectedDraftTokens[bi]; - auto const slot = batchSlots[bi]; - for (SizeType32 ti = 0; ti < expectedDraftTokensBatch.size(); ++ti) - { - EXPECT_EQ(expectedDraftTokensBatch[ti], nextDraftTokens[slot * (mMaxDecodingTokens - 1) + ti]) - << "bi " << bi << " ti " << ti; - } - } - for (SizeType32 bi = 0; bi < mBatchSize; ++bi) - { - auto const slot = batchSlots[bi]; - EXPECT_EQ(finished[bi], finishedPtr[slot].isFinished()); - } -} - -template -void MedusaDecodingLayerTest::runTest(std::vector>> const& expectedOutTokens, - std::vector> const& expectedDraftTokens, std::vector const& finished, - SamplingParams& params) -{ - mBatchSize = params.batchSize; - mMaxBatchSize = 2 * mBatchSize; - - allocateBuffers(); - - setup(params); - - auto inputTensors = createInputTensors(); - auto outputTensors = createOutputTensors(); - - mDecodingWorkspace->setDeviceBatchSlots(mBatchSlots); - mMedusaDecodingLayer->forwardAsync(outputTensors, inputTensors, mDecodingWorkspace); - - mStream->synchronize(); - - checkResult(expectedOutTokens, expectedDraftTokens, finished, params); -} - -template class MedusaDecodingLayerTest; -template class MedusaDecodingLayerTest; - -TYPED_TEST_SUITE(MedusaDecodingLayerTest, FloatAndHalfTypes); - -TYPED_TEST(MedusaDecodingLayerTest, SimpleTestBS1) -{ - SamplingParams params; - params.runtimeTopK = {1}; - params.runtimeHeadsTopK = {{2, 3, 2, 1}}; - params.draftIds = {{4, 0, 2, 1, 3, 4, 3, 0, 2, 3, 4}}; - params.paths = {{0, 1, 2, 3, -1}}; - params.treeIds = {{0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2}}; - params.tokensPerStep = {12}; - params.acceptedCumSum = {0, 3}; - params.packedPaths = {0, 1, 2}; - params.batchSize = 1; - - std::vector>> expectedOutTokens = {{{4}, {0}, {2}, {1}}}; - std::vector> expectedDraftTokens = {{1, 2, 1, 2, 3, 2, 3, 0, 1, 2, 1}}; - std::vector finished = {false}; - this->runTest(expectedOutTokens, expectedDraftTokens, finished, params); -} - -TYPED_TEST(MedusaDecodingLayerTest, SimpleTestBS4) -{ - // Target Ids to be sampled - // [4, 0, 2, 1, 3, 4, 3, 0, 2, 3, 4, 1] - SamplingParams params; - params.runtimeTopK = {1, 1, 1, 1}; - params.runtimeHeadsTopK = {{2, 3, 2, 1}, {1, 2, 3, 4}, {3, 1, 1, 1}, {1, 1, 1, 1}}; - // clang-format off - params.draftIds = {{4, 0, 2, 1, 3, 4, 4, 0, 2, 3, 4}, - {4, 0, 2, 1, 4, 4, 4, 0, 2, 2, 4}, - {4, 0, 4, 1, 1, 4, 4, 0, 2, 0, 4}, - {4, 0, 2, 1, 3, 2, 4, 0, 2, 3, 4}}; - params.paths = {{0, 7, 2, 8, -1, - 0, 3, -1, -1, -1}, - {0, 5, 7, 8, 10, - 0, 3, -1, -1, -1}, - {0, 8, 2, 9, -1, - 0, 3, 5, 6, -1, - 0, 3, 5, 7, 10}, - {0, 1, 2, 6, -1, - 0, 3, -1, -1, -1}}; - params.treeIds = {{0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2}, - {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, - {0, 1, 2, 3, 4, 5, 0, 1, 2, 3}, - {0, 1, 2, 3, 1, 2}}; - params.acceptedCumSum = {0, 2, 6, 10, 13}; - params.packedPaths = {6, 1, 4, 6, 7, 9, 2, 4, 6, 9, 0, 1, 5}; - // clang-format on - params.tokensPerStep = {12, 11, 11, 7}; - params.batchSize = 4; - - std::vector>> expectedOutTokens - = {{{4}, {0}, {2}}, {{4}, {4}, {0}, {2}, {4}}, {{4}, {1}, {4}, {0}, {4}}, {{4}, {0}, {2}, {3}}}; - std::vector> expectedDraftTokens = {{2, 3, 4, 5, 6, 1, 2, 1, 2, 3, 4}, - {4, 1, 2, 4, 5, 6, 0, 1, 2, 3}, {4, 5, 6, 1, 4, 0, 4, 5, 6, 1}, {3, 3, 0, 3, 3, 0}}; - std::vector finished = {false, false, false, false}; - this->runTest(expectedOutTokens, expectedDraftTokens, finished, params); -} - -TYPED_TEST(MedusaDecodingLayerTest, SimpleTestEndIdNotSelected) -{ - // Target Ids to be sampled - // [4, 0, 2, 1, 3, 4, 3, 0, 2, 3, 4, 1] - SamplingParams params; - params.runtimeTopK = {1}; - params.runtimeHeadsTopK = {{1, 1, 1, 1}}; - params.draftIds = {{4, 0, 4, 1, 3, 2, 3, 0, 2, 3, 4}}; - // clang-format off - params.paths = {{0, 3, 4, 5, -1, - 0, 1, 2, 6, -1}}; - params.treeIds = {{0, 1, 2, 3, 0, 1, 2, 3, 3, 2, 1}}; - // clang-format on - params.tokensPerStep = {12}; - params.acceptedCumSum = {0, 3}; - params.packedPaths = {0, 1, 5}; - params.batchSize = 1; - params.endId = 1; - - std::vector>> expectedOutTokens = {{{4}, {0}, {2}, {3}}}; - std::vector> expectedDraftTokens = {{3, 3, 0, 3, 3, 3, 0, 3, 3, 0, 3}}; - std::vector finished = {false}; - this->runTest(expectedOutTokens, expectedDraftTokens, finished, params); -} - -TYPED_TEST(MedusaDecodingLayerTest, SimpleTestEndIdSelected) -{ - // Target Ids to be sampled - // [4, 0, 2, 1, 3, 4, 3, 0, 2, 3, 4, 1] - SamplingParams params; - params.runtimeTopK = {1}; - params.runtimeHeadsTopK = {{1, 1, 1, 1}}; - params.draftIds = {{4, 0, 4, 1, 3, 2, 3, 0, 2, 3, 4}}; - // clang-format off - params.paths = {{0, 3, 4, 5, -1, - 0, 11, 7, 9, -1}}; - params.treeIds = {{0, 1, 2, 3, 0, 1, 2, 3, 3, 2, 1}}; - // clang-format on - params.tokensPerStep = {12}; - params.acceptedCumSum = {0, 0}; - params.packedPaths = {}; - params.batchSize = 1; - params.endId = 1; - - std::vector>> expectedOutTokens = {{{4}}}; - std::vector> expectedDraftTokens = {{1, 1, 2, 0, 1, 1, 2, 0, 0, 2, 1}}; - std::vector finished = {true}; - this->runTest(expectedOutTokens, expectedDraftTokens, finished, params); -} -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.h b/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.h deleted file mode 100644 index 673f52ff6dde..000000000000 --- a/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.h +++ /dev/null @@ -1,117 +0,0 @@ -/* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include - -#include - -#include "tensorrt_llm/layers/medusaDecodingLayer.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" - -namespace tensorrt_llm::tests::layers -{ - -struct SamplingParams -{ - tensorrt_llm::runtime::SizeType32 batchSize; - std::vector runtimeTopK; - std::vector> runtimeHeadsTopK; - std::vector> draftIds; - std::vector> paths; - std::vector> treeIds; - std::vector tokensPerStep; - std::vector acceptedCumSum; - std::vector packedPaths; - std::optional endId; -}; - -template -class MedusaDecodingLayerTest : public testing::Test -{ -private: - void SetUp() override; - -public: - using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; - using BufferPtr = tensorrt_llm::runtime::IBuffer::SharedPtr; - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TokenIdType = tensorrt_llm::runtime::TokenIdType; - -private: - SizeType32 mBatchSize{6}; - SizeType32 mMaxBatchSize{2 * mBatchSize}; - SizeType32 const mVocabSize{9}; - SizeType32 const mVocabSizePadded{mVocabSize}; - SizeType32 const mMaxDecodingTokens{12}; - SizeType32 const mMaxDraftPathLen{4}; - - SizeType32 const mMaxSeqLen{mMaxDecodingTokens}; - TokenIdType mEndId{mVocabSize}; - - bool mUseLogitsVec{false}; - - TensorPtr mTargetLogitsDevice; - TensorPtr mMedusaLogitsDevice; - - TensorPtr mFinishedDevice; - TensorPtr mSeqLengthsDevice; - TensorPtr mAcceptedLengths; - TensorPtr mOutputIdsDevice; - TensorPtr mNextDraftTokensDevice; - - TensorPtr mPathsDevice; - TensorPtr mTreeIdsDevice; - TensorPtr mAcceptedLengthCumSumDevice; - TensorPtr mPackedPathsDevice; - TensorPtr mEndIdsDevice; - TensorPtr mBatchSlots; - - TensorPtr mTokensPerStepDevice; - - std::vector mLogitsVec; - - std::shared_ptr mStream; - std::shared_ptr mBufferManager; - std::shared_ptr> mMedusaDecodingLayer; - std::shared_ptr mDecodingWorkspace; - -private: - void allocateBuffers(); - - void setup(SamplingParams& params); - - std::shared_ptr createInputTensors(); - - std::shared_ptr createOutputTensors(); - - void checkResult(std::vector>> const& expectedOutTokens, - std::vector> const& expectedDraftTokens, std::vector const& finished, - SamplingParams& params); - -public: - void runTest(std::vector>> const& expectedOutTokens, - std::vector> const& expectedDraftTokens, std::vector const& finished, - SamplingParams& params); -}; - -typedef testing::Types FloatAndHalfTypes; - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/randomLlm.cpp b/cpp/tests/unit_tests/layers/randomLlm.cpp deleted file mode 100644 index 63aa85eaad16..000000000000 --- a/cpp/tests/unit_tests/layers/randomLlm.cpp +++ /dev/null @@ -1,338 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "tests/unit_tests/layers/randomLlm.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::tests::layers -{ - -using namespace tensorrt_llm::layers; - -TensorPtr initTensor(std::string str, std::optional shape) -{ - auto shape1d = ITensor::makeShape({static_cast(str.size())}); - if (shape) - { - TLLM_CHECK(ITensor::volume(shape1d) == ITensor::volume(shape.value())); - } - TensorPtr tensor = BufferManager::cpu(shape.value_or(shape1d), tensorrt_llm::DataType::kINT32); - auto tensorRange = BufferRange(*tensor); - std::copy(str.begin(), str.end(), tensorRange.begin()); - return tensor; -} - -TensorConstPtr RandomTokenLogits::tokenToLogits(TokenIdType token) const -{ - TensorPtr logits = BufferManager::cpu(mVocabulary->getShape(), tensorrt_llm::DataType::kFLOAT); - tokenToLogits(logits, token); - return logits; -} - -void RandomTokenLogits::tokenToLogits(TensorPtr const& logits, TokenIdType token) const -{ - TLLM_CHECK_WITH_INFO(logits->shapeEquals({getVocabSize()}), "%s != {%d}", - ITensor::toString(logits->getShape()).c_str(), getVocabSize()); - - auto logitsRange = BufferRange(*logits); - auto vocabRange = BufferRange(*mVocabulary); - auto itl = logitsRange.begin(); - auto itv = vocabRange.begin(); - for (; itl != logitsRange.end() && itv != vocabRange.end(); itl++, itv++) - { - bool match = (*itv == token); - *itl = (match ? 1.0 : 0.0) + (static_cast(rand() % 256) / 1000.0); - } -} - -TokenIdType RandomTokenLogits::logitsToToken(TensorConstPtr const& logits) const -{ - TLLM_CHECK(logits->shapeEquals({getVocabSize()})); - auto logitsRange = BufferRange(*logits); - auto vocabRange = BufferRange(*mVocabulary); - float max = -FLT_MAX; - TokenIdType result; - auto itl = logitsRange.begin(); - auto itv = vocabRange.begin(); - for (; itl != logitsRange.end() && itv != vocabRange.end(); itl++, itv++) - { - float cur = exp(*itl); - if (cur > max) - { - max = cur; - result = *itv; - } - } - return result; -} - -std::list RandomTokenLogits::stringToLogits(std::string tokens) const -{ - std::list result; - for (auto& token : tokens) - { - result.push_back(tokenToLogits(static_cast(token))); - } - return result; -} - -void RandomTokenLogits::stringToLogits(TensorPtr const& logits, std::string tokens) const -{ - TLLM_CHECK(logits->shapeEquals({static_cast(tokens.size()), getVocabSize()})); - - auto i = 0; - for (auto& token : tokens) - { - tokenToLogits(ITensor::at(logits, {i++}), static_cast(token)); - } -} - -void RandomTokenLogits::tensorToLogits(TensorPtr const& logits, TensorConstPtr const& tokens) const -{ - TLLM_CHECK(ITensor::volume(logits->getShape()) == ITensor::volume(tokens->getShape()) * getVocabSize()); - // TLLM_CHECK(logits->shapeEquals({static_cast(tokens.size()), getVocabSize()})); - auto tokensRange = BufferRange(*tokens); - auto i = 0; - for (auto it = tokensRange.begin(); it != tokensRange.end(); it++) - { - tokenToLogits(ITensor::at(logits, {i++}), *it); - } -} - -std::string RandomTokenLogits::logitsToString(std::list logits) const -{ - std::string result; - for (auto& token : logits) - { - result.push_back(logitsToToken(token)); - } - return result; -} - -std::string RandomTokenLogits::logitsToString(TensorConstPtr const& logits) const -{ - auto len = logits->getShape().d[0]; - std::string result; - for (auto i = 0; i < len; i++) - { - result.push_back(logitsToToken(ITensor::at(logits, {i}))); - } - return result; -} - -void RandomTokenLogits::logitsToTensor(TensorPtr const& tokens, TensorConstPtr const& logits) const -{ - auto len = logits->getShape().d[0]; - TLLM_CHECK(tokens->getShape().d[0] >= len); - auto tokensRange = BufferRange(*tokens); - for (auto i = 0; i < len; i++) - { - tokensRange[i] = logitsToToken(ITensor::at(logits, {i})); - } -} - -TensorConstPtr RandomTokenLogits::logitsToTensor(TensorConstPtr const& logits) const -{ - auto len = logits->getShape().d[0]; - TensorPtr result = BufferManager::cpu(ITensor::makeShape({len}), tensorrt_llm::DataType::kINT32); - logitsToTensor(result, logits); - return result; -} - -SizeType32 RandomTokenLogits::getVocabSize() const -{ - return ITensor::volume(mVocabulary->getShape()); -} - -TokenIdType const RandomTokenLogits::getInvalidToken() const -{ - return *(BufferRange(*mVocabulary).end() - 1); -} - -TokenIdType const RandomTokenLogits::getEndToken() const -{ - return *(BufferRange(*mVocabulary).end() - 2); -} - -void RandomLlm::sampleByMask(TensorPtr const& inout, TensorConstPtr const& mask) const -{ - auto len = ITensor::volume(mask->getShape()); - TLLM_CHECK(len == ITensor::volume(mask->getShape())); - auto inoutRange = BufferRange(*inout); - auto maskRange = BufferRange(*mask); - auto invalid = mTable->getInvalidToken(); - - for (SizeType32 i = 0; i < len; i++) - { - if (!maskRange[i]) - { - inoutRange[i] = invalid; - } - } -} - -bool RandomLlm::verify(SizeType32 const offset, TensorConstPtr const& script) const -{ - auto oracleRange = BufferRange(*mOracle); - auto scriptRange = BufferRange(*script); - auto len = ITensor::volume(script->getShape()); - auto result = std::equal(oracleRange.begin() + offset, oracleRange.begin() + offset + len, scriptRange.begin()); - if (!result) - { - std::string gold(len, '#'); - std::string wrong(len, '#'); - std::copy(oracleRange.begin() + offset, oracleRange.begin() + offset + len, gold.begin()); - std::copy(scriptRange.begin(), scriptRange.end(), wrong.begin()); - TLLM_CHECK_WITH_INFO(result, "len=%ld, gold='%s', script='%s'", len, gold.c_str(), wrong.c_str()); - } - return result; -} - -void RandomLlm::forward(TensorPtr const& output, runtime::SizeType32 startId, TensorConstPtr const& input, - TensorConstPtr const& offsets, TensorConstPtr const mask) const -{ - TensorPtr posIds = BufferManager::cpu(input->getShape(), tensorrt_llm::DataType::kINT32); - BufferRange idRange(*posIds); - BufferRange offsetRange(*offsets); - for (auto i = 0; i < idRange.size(); i++) - { - idRange[i] = startId + offsetRange[i]; - } - forward(output, input, posIds, mask); -} - -void RandomLlm::forward(TensorPtr const& output, TensorConstPtr const& input, TensorConstPtr const& position, - TensorConstPtr const mask) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(ITensor::volume(input->getShape()) == ITensor::volume(position->getShape())); - TLLM_CHECK(ITensor::volume(output->getShape()) == ITensor::volume(input->getShape()) * mTable->getVocabSize()); - - TensorPtr tokens = BufferManager::cpu(input->getShape(), tensorrt_llm::DataType::kINT32); - foretell(tokens, input, position, mask); - // foretellOld(tokens, input, position); - mTable->tensorToLogits(output, tokens); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRandomLlm::foretell(TensorPtr const& output, TensorConstPtr const& input, TensorConstPtr const& position, - TensorConstPtr const mask) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto len = ITensor::volume(input->getShape()); - TLLM_CHECK(ITensor::volume(position->getShape()) == len); - TLLM_CHECK(ITensor::volume(output->getShape()) >= len); - if (mask) - { - TLLM_CHECK(ITensor::volume(mask->getShape()) >= len * len); - TLLM_CHECK(mask->getShape().d[0] >= len); - TLLM_CHECK(mask->getShape().d[1] >= len); - } - - TensorPtr maskRebuilt = BufferManager::cpu(ITensor::makeShape({len, len}), tensorrt_llm::DataType::kBOOL); - posIdsToMask(maskRebuilt, position); - - auto outputRange = BufferRange(*output); - auto inputRange = BufferRange(*input); - auto positionRange = BufferRange(*position); - auto maskLocation = mask ? BufferLocation(*mask) : BufferLocation(*maskRebuilt); - auto oracleRange = BufferRange(*mOracle); - auto olen = ITensor::volume(mOracle->getShape()); - - auto verifyStart = 2; - for (; verifyStart < len - 1; verifyStart++) - { - if (positionRange[verifyStart] == positionRange[0] + 1) - { - break; - } - } - - auto invalid = mTable->getInvalidToken(); - TLLM_CHECK(positionRange[0] + 1 < olen); - for (auto i = 0; i < len; i++) - { - bool legal = positionRange[i] + 1 < olen; - bool right = true; - for (auto j = 0; j < i; j++) - { - right &= maskLocation.at(i, j) ? oracleRange[positionRange[j]] == inputRange[j] : true; - } - if (i < verifyStart && false) - { // lookahead might be right. Since we verify lookahead branch, then must be right. - outputRange[i] = ((right || rand() % 5) && legal) ? oracleRange[positionRange[i] + 1] : invalid; - } - else - { // verify should be wrong. - outputRange[i] = (right && legal) ? oracleRange[positionRange[i] + 1] : invalid; - } - } -} - -void LookaheadRandomLlm::posIdsToMask(TensorPtr const& mask, TensorConstPtr const& posIds) const -{ - auto len = ITensor::volume(posIds->getShape()); - TLLM_CHECK(ITensor::volume(mask->getShape()) >= len * len); - auto posIdsRange = BufferRange(*posIds); - auto maskRange = BufferRange(*mask); - - for (auto i = 0; i < maskRange.size(); i++) - { - maskRange[i] = false; - } - - std::vector> stack; - stack.push_back(std::make_pair(0, posIdsRange[0])); - maskRange[0 * len + 0] = true; - for (auto i = 1; i < len; i++) - { - auto cur = posIdsRange[i]; - while (stack.size() > 0 && cur <= stack.back().second) - { - stack.pop_back(); - } - TLLM_CHECK(stack.size() > 0 ? cur == stack.back().second + 1 : true); - stack.push_back(std::make_pair(i, cur)); - for (auto prev : stack) - { - maskRange[i * len + prev.first] = true; - } - } -} - -void LookaheadRandomLlm::maskToPosIds(TensorPtr const& posIds, TensorConstPtr const& mask, SizeType32 start) const -{ - auto len = ITensor::volume(posIds->getShape()); - TLLM_CHECK(ITensor::volume(mask->getShape()) >= len * len); - auto posIdsRange = BufferRange(*posIds); - auto maskLocation = BufferLocation(*mask); - for (auto i = 0; i < len; i++) - { - posIdsRange[i] = start; - for (auto j = 0; j < i; j++) - { - posIdsRange[i] += maskLocation.at(i, j); - } - } -} - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/randomLlm.h b/cpp/tests/unit_tests/layers/randomLlm.h deleted file mode 100644 index b0f0564944dc..000000000000 --- a/cpp/tests/unit_tests/layers/randomLlm.h +++ /dev/null @@ -1,152 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" - -namespace tensorrt_llm::tests::layers -{ -using namespace tensorrt_llm::runtime; -using TensorPtr = runtime::ITensor::SharedPtr; -using TensorConstPtr = runtime::ITensor::SharedConstPtr; - -//! Initialize a tensor with data from string @param str. Shape {str.size} by default. -TensorPtr initTensor(std::string str, std::optional shape = std::nullopt); - -//! Convert tokens to logits and vice versa according to a vocabulary. -class RandomTokenLogits -{ -public: - RandomTokenLogits(TensorConstPtr const& vocab) - : mVocabulary(vocab) - { - } - - RandomTokenLogits(std::string vocab) - : mVocabulary(initTensor(vocab)) - { - } - - TensorConstPtr tokenToLogits(TokenIdType token) const; - void tokenToLogits(TensorPtr const& logits, TokenIdType token) const; - - TokenIdType logitsToToken(TensorConstPtr const& logits) const; - - std::list stringToLogits(std::string tokens) const; - void stringToLogits(TensorPtr const& logits, std::string tokens) const; - void tensorToLogits(TensorPtr const& logits, TensorConstPtr const& tokens) const; - - std::string logitsToString(std::list logits) const; - std::string logitsToString(TensorConstPtr const& logits) const; - TensorConstPtr logitsToTensor(TensorConstPtr const& logits) const; - void logitsToTensor(TensorPtr const& tokens, TensorConstPtr const& logits) const; - - SizeType32 getVocabSize() const; - //! @return the last token in mVocabulary as invalid token; - virtual TokenIdType const getInvalidToken() const; - //! @return the second-to-last token in mVocabulary as end token; - virtual TokenIdType const getEndToken() const; - -private: - TensorConstPtr const mVocabulary; -}; - -//! vocabulary is ascii table from 0 to 127. tokenId == token. -class AsciiRandomTokenLogits : public RandomTokenLogits -{ -public: - AsciiRandomTokenLogits() - : RandomTokenLogits( - []() - { - auto vocab = BufferManager::cpu(ITensor::makeShape({128}), tensorrt_llm::DataType::kINT32); - auto vocabRange = BufferRange(*vocab); - TokenIdType token{0}; - std::for_each(vocabRange.begin(), vocabRange.end(), [&token](auto& v) { v = token++; }); - return vocab; - }()) - { - } - - virtual TokenIdType const getInvalidToken() const - { - return static_cast('#'); - } - - virtual TokenIdType const getEndToken() const - { - return static_cast('&'); - } -}; - -//! random LLM to simulate functions of a real LLM. -class RandomLlm -{ -public: - RandomLlm(std::shared_ptr const table, std::string oracle, runtime::SizeType32 id = 0) - : mTable(table) - , mOracle(initTensor(oracle)) - , mId(id) - { - } - - // simulate forward in a LLM. - void forward(TensorPtr const& output, runtime::SizeType32 startId, TensorConstPtr const& input, - TensorConstPtr const& offsets, TensorConstPtr const mask = nullptr) const; - void forward(TensorPtr const& output, TensorConstPtr const& input, TensorConstPtr const& position, - TensorConstPtr const mask = nullptr) const; - //! set inout[i] invalid if mask[i]==false; - void sampleByMask(TensorPtr const& inout, TensorConstPtr const& mask) const; - //! @return true when @param script is a sub-string started from @param offset. - bool verify(SizeType32 const offset, TensorConstPtr const& script) const; - - //! foretell @param output tokens from @param input tokens and @param position ids. - //! It depends on different algorithms implementations. - virtual void foretell(TensorPtr const& output, TensorConstPtr const& input, TensorConstPtr const& position, - TensorConstPtr const mask = nullptr) const - = 0; - -protected: - std::shared_ptr const mTable; - TensorConstPtr const mOracle; - runtime::SizeType32 const mId; -}; - -//! a lookahead implementation for RandomLlm. -class LookaheadRandomLlm : public RandomLlm -{ -public: - LookaheadRandomLlm( - std::shared_ptr const table, std::string oracle, runtime::SizeType32 id = 0) - : RandomLlm(table, oracle, id) - { - } - - void foretell(TensorPtr const& output, TensorConstPtr const& input, TensorConstPtr const& position, - TensorConstPtr const mask = nullptr) const override; - -private: - void posIdsToMask(TensorPtr const& mask, TensorConstPtr const& posIds) const; - void maskToPosIds(TensorPtr const& posIds, TensorConstPtr const& mask, runtime::SizeType32 start) const; -}; - -} // namespace tensorrt_llm::tests::layers diff --git a/cpp/tests/unit_tests/layers/samplingLayerTest.cpp b/cpp/tests/unit_tests/layers/samplingLayerTest.cpp deleted file mode 100644 index 7960e4eca314..000000000000 --- a/cpp/tests/unit_tests/layers/samplingLayerTest.cpp +++ /dev/null @@ -1,376 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/executor/types.h" -#include "tests/unit_tests/layers/baseSamplingLayerTest.h" - -namespace -{ - -namespace tle = tensorrt_llm::executor; - -using namespace tensorrt_llm::tests::layers::sampling; -using namespace tensorrt_llm::runtime; - -template -class SamplingLayerTest : public BaseSamplingLayerTest -{ - void SetUp() override - { - this->mStream = std::make_shared(); - this->mBufferManager = std::make_shared(this->mStream); - } - - void initLayer(TestSamplingParams const& params) override - { - auto decodingMode = tle::DecodingMode::Auto(); - if (params.topKs.size() && params.topPs.size()) - { - decodingMode = tle::DecodingMode::TopKTopP(); - } - else if (params.topKs.size()) - { - decodingMode = tle::DecodingMode::TopK(); - } - else if (params.topPs.size()) - { - decodingMode = tle::DecodingMode::TopP(); - } - - auto const decodingDomain - = tensorrt_llm::layers::DecoderDomain(this->maxBatchSize(), 1, this->mVocabSize, this->mVocabSizePadded); - this->mSamplingLayer = std::make_shared>( - decodingMode, decodingDomain, this->mBufferManager); - } -}; - -TYPED_TEST_SUITE(SamplingLayerTest, FloatAndHalfTypes); - -TYPED_TEST(SamplingLayerTest, TopKToPPSkipDecode) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, TopKSkipDecodeTopP) -{ - SizeType32 topK = 0; - float topP = 0.5f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, BatchTopKTopP) -{ - std::vector topKs = {0, 2, 1, 0, 1, 0}; - std::vector topPs = {0.3f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4}, {4, 5}, {4}, {4, 5}, {4}, {4, 5}, // step 0 - {0}, {0, 1}, {0}, {0, 1}, {0}, {0, 1}, // step 1 - {2}, {2, 3}, {2}, {2, 3}, {2}, {2, 3}, // step 2 - {0}, {0, 1}, {0}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, TopPDecay) -{ - TestSamplingParams params; - params.topPs = {0.8f, 0.5f, 0.3f, 0.2f, 0.5f, 1.0f}; - params.decay = {0.3f, 0.3f, 0.3f, 0.9f, 0.3f, 0.8f}; - params.topPResetIds = {2, -1, 2, -1, 2, -1}; - params.minTopP = {0.5f, 0.1f, 0.3f, 0.1f, 0.1f, 0.1f}; - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4, 5}, {4}, {4}, {4, 5}, {4, 5, 6, 7}, // step 0 - {0, 1}, {0}, {0}, {0}, {0}, {0, 1, 2}, // step 1 - {2, 3}, {2}, {2}, {2}, {2}, {2, 3}, // step 2 - {0, 1, 2}, {0}, {0}, {0}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, TopK) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, TopK1TopP0) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, BatchTopK) -{ - std::vector topKs = {2, 1, 1, 2, 1, 1}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {2, 3}, {2}, {2}, {2, 3}, {2}, {2}, // step 2 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, TopKTopP) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, BatchTopKTopP1) -{ - std::vector topKs = {2, 2, 1, 2, 2, 1}; - float topP = 0.3; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, BatchTopKBatchTopP) -{ - std::vector topKs = {2, 2, 0, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4, 5}, {4, 5}, {4}, {4}, // step 0 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0}, // step 1 - {2, 3}, {2}, {2, 3}, {2, 3}, {2}, {2}, // step 2 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsZeroTopK) -{ - SizeType32 topK = 0; - TestSamplingParams params; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsZeroTopP) -{ - float topP = 0; - SizeType32 topK = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsZeroTopKTopP) -{ - SizeType32 topK = 0; - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsZeroBatchTopKTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsZeroTopKBatchTopP) -{ - SizeType32 topK = 0; - std::vector topPs = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - TestSamplingParams params; - params.topPs = topPs; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsBatchTopKContainZero) -{ - std::vector topKs = {2, 1, 0, 0, 2, 1}; - TestSamplingParams params; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, InvalidArgsBatchTopKTopPContainZero) -{ - std::vector topKs = {2, 2, 1, 0, 2, 0}; - float topP = 0.0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2, 3}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, OnlyTopK) -{ - std::vector topKs = {2, 2, 1, 0, 2, 0}; - TestSamplingParams params; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2, 3}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(SamplingLayerTest, OnlyTopP) -{ - std::vector topPs = {0.3f, 0.3f, 0.3f, 0.3f, 0.3f, 0.3f}; - TestSamplingParams params; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -} // namespace diff --git a/cpp/tests/unit_tests/layers/topKSamplingLayerTest.cpp b/cpp/tests/unit_tests/layers/topKSamplingLayerTest.cpp deleted file mode 100644 index 73d8f881c304..000000000000 --- a/cpp/tests/unit_tests/layers/topKSamplingLayerTest.cpp +++ /dev/null @@ -1,309 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tests/unit_tests/layers/baseSamplingLayerTest.h" - -namespace -{ - -using namespace tensorrt_llm::tests::layers::sampling; -using namespace tensorrt_llm::runtime; - -template -class TopKSamplingLayerTest : public BaseSamplingLayerTest -{ - void SetUp() override - { - this->mStream = std::make_shared(); - this->mBufferManager = std::make_shared(this->mStream); - } - - void initLayer(TestSamplingParams const& params) override - { - auto const decodingDomain - = tensorrt_llm::layers::DecoderDomain(this->maxBatchSize(), 1, this->mVocabSize, this->mVocabSizePadded); - this->mSamplingLayer - = std::make_shared>(decodingDomain, this->mBufferManager); - } -}; - -TYPED_TEST_SUITE(TopKSamplingLayerTest, FloatAndHalfTypes); - -TYPED_TEST(TopKSamplingLayerTest, TopK) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, {4, 5}, // step 0 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, // step 1 - {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, {2, 3}, // step 2 - {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, TopK1TopP0) -{ - SizeType32 topK = 1; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, BatchTopK) -{ - std::vector topKs = {2, 1, 1, 2, 1, 1}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = {1.0f}; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {2, 3}, {2}, {2}, {2, 3}, {2}, {2}, // step 2 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, SkipDecode) -{ - // Skip topK decode - float topP = 0.3; - TestSamplingParams params; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {0}, {0}, {0}, {0}, {0}, {0}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {0}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, TopKTopP) -{ - SizeType32 topK = 2; - float topP = 0.3; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, BatchTopKTopP) -{ - std::vector topKs = {2, 2, 1, 2, 2, 1}; - float topP = 0.3; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, TopKBatchTopP) -{ - SizeType32 topK = 2; - std::vector topPs = {0.5, 0.3, 0.5, 0.5, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4, 5}, {4, 5}, {4}, {4, 5}, // step 0 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0, 1}, // step 1 - {2, 3}, {2}, {2, 3}, {2, 3}, {2}, {2, 3}, // step 2 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, BatchTopKBatchTopP) -{ - std::vector topKs = {2, 2, 1, 2, 2, 1}; - std::vector topPs = {0.0, 0.3, 0.5, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4, 5}, {4}, {4}, // step 0 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {2, 3}, {2}, {2}, {2, 3}, {2}, {2}, // step 2 - {0, 1}, {0}, {0}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsZeroTopK) -{ - SizeType32 topK = 0; - TestSamplingParams params; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsZeroTopP) -{ - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsZeroTopKTopP) -{ - SizeType32 topK = 0; - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsZeroBatchTopKTopP) -{ - std::vector topKs = {0, 0, 0, 0, 0, 0}; - float topP = 0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsZeroTopKBatchTopP) -{ - SizeType32 topK = 0; - std::vector topPs = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - TestSamplingParams params; - params.topPs = topPs; - params.topKs = {topK}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsBatchTopKContainZero) -{ - std::vector topKs = {2, 1, 0, 0, 2, 1}; - TestSamplingParams params; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsBatchTopKTopPContainZero) -{ - std::vector topKs = {2, 2, 1, 0, 2, 0}; - float topP = 0.0; - TestSamplingParams params; - params.topPs = {topP}; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4, 5}, {4}, {4}, {4, 5}, {4}, // step 0 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0}, // step 1 - {2, 3}, {2, 3}, {2}, {2}, {2, 3}, {2}, // step 2 - {0, 1}, {0, 1}, {0}, {0}, {0, 1}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopKSamplingLayerTest, InvalidArgsBatchTopKBatchTopPContainZero) -{ - std::vector topKs = {0, 2, 1, 2, 2, 0}; - std::vector topPs = {0.0, 0.3, 0.9, 0.0, 0.3, 0.5}; - TestSamplingParams params; - params.topPs = topPs; - params.topKs = topKs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4, 5}, {4}, {0}, // step 0 - {0}, {0}, {0}, {0, 1}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2, 3}, {2}, {0}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -} // namespace diff --git a/cpp/tests/unit_tests/layers/topPSamplingLayerTest.cpp b/cpp/tests/unit_tests/layers/topPSamplingLayerTest.cpp deleted file mode 100644 index 89d3428e9dd5..000000000000 --- a/cpp/tests/unit_tests/layers/topPSamplingLayerTest.cpp +++ /dev/null @@ -1,193 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tests/unit_tests/layers/baseSamplingLayerTest.h" - -namespace -{ - -using namespace tensorrt_llm::tests::layers::sampling; -using namespace tensorrt_llm::runtime; - -template -class TopPSamplingLayerTest : public BaseSamplingLayerTest -{ - void SetUp() override - { - this->mStream = std::make_shared(); - this->mBufferManager = std::make_shared(this->mStream); - - int device; - cudaGetDevice(&device); - cudaGetDeviceProperties(&mDeviceProp, device); - - this->mComputeProbs = true; - } - - void initLayer(TestSamplingParams const& params) override - { - auto const decodingDomain - = tensorrt_llm::layers::DecoderDomain(this->maxBatchSize(), 1, this->mVocabSize, this->mVocabSizePadded); - this->mSamplingLayer = std::make_shared>( - decodingDomain, this->mBufferManager, &mDeviceProp); - } - -protected: - cudaDeviceProp mDeviceProp{}; -}; - -TYPED_TEST_SUITE(TopPSamplingLayerTest, FloatAndHalfTypes); - -TYPED_TEST(TopPSamplingLayerTest, TopKSkipDecode) -{ - SizeType32 topK = 2; - float topP = 0.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {0}, {0}, {0}, {0}, {0}, {0}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {0}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, TopKTopPSkipDecode) -{ - SizeType32 topK = 2; - float topP = 1.0f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {0}, {0}, {0}, {0}, {0}, {0}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {0}, {0}, {0}, {0}, {0}, {0}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, BatchTopKTopP) -{ - std::vector topKs = {0, 1, 1, 0, 1, 0}; - std::vector topPs = {0.3f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f}; - TestSamplingParams params; - params.topKs = topKs; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4}, {0}, {0}, {4, 5}, {0}, {4, 5}, // step 0 - {0}, {0}, {0}, {0, 1}, {0}, {0, 1}, // step 1 - {2}, {0}, {0}, {2, 3}, {0}, {2, 3}, // step 2 - {0}, {0}, {0}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, TopP) -{ - SizeType32 topK = 0; - float topP = 0.3f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4}, {4}, {4}, {4}, // step 0 - {0}, {0}, {0}, {0}, {0}, {0}, // step 1 - {2}, {2}, {2}, {2}, {2}, {2}, // step 2 - {0}, {0}, {0}, {0}, {0}, {0} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, BatchTopP) -{ - std::vector topPs = {0.3f, 0.3f, 0.5f, 0.8f, 0.5f, 0.8f}; - TestSamplingParams params; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4}, {4}, {4, 5}, {4, 5, 6}, {4, 5}, {4, 5, 6}, // step 0 - {0}, {0}, {0, 1}, {0, 1, 2}, {0, 1}, {0, 1, 2}, // step 1 - {2}, {2}, {2, 3}, {2, 3, 4}, {2, 3}, {2, 3, 4}, // step 2 - {0}, {0}, {0, 1}, {0, 1, 2}, {0, 1}, {0, 1, 2} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, TopKBatchTopP) -{ - std::vector topPs = {0.5f, 0.3f, 0.5f, 0.5f, 0.3f, 0.5f}; - TestSamplingParams params; - params.topPs = topPs; - std::vector> expectedOutputIds{ - // batch - {4, 5}, {4}, {4, 5}, {4, 5}, {4}, {4, 5}, // step 0 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0, 1}, // step 1 - {2, 3}, {2}, {2, 3}, {2, 3}, {2}, {2, 3}, // step 2 - {0, 1}, {0}, {0, 1}, {0, 1}, {0}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, TopPDecay) -{ - TestSamplingParams params; - params.topPs = {0.8f, 0.5f, 0.3f, 0.2f, 0.5f, 1.0f}; - params.decay = {0.3f, 0.3f, 0.3f, 0.9f, 0.3f, 0.8f}; - params.topPResetIds = {2, -1, 2, -1, 2, -1}; - params.minTopP = {0.5f, 0.1f, 0.3f, 0.1f, 0.1f, 0.1f}; - std::vector> expectedOutputIds{ - // batch - {4, 5, 6}, {4, 5}, {4}, {4}, {4, 5}, {4, 5, 6, 7}, // step 0 - {0, 1}, {0}, {0}, {0}, {0}, {0, 1, 2}, // step 1 - {2, 3}, {2}, {2}, {2}, {2}, {2, 3}, // step 2 - {0, 1, 2}, {0}, {0}, {0}, {0, 1}, {0, 1} // step 3 - }; - this->runTest(expectedOutputIds, params); -} - -TYPED_TEST(TopPSamplingLayerTest, LargeBatch) -{ - SizeType32 topK = 0; - float topP = 0.3f; - TestSamplingParams params; - params.topKs = {topK}; - params.topPs = {topP}; - - // Force to use more than 1 block - params.batchSize = this->mDeviceProp.maxThreadsPerBlock + 1; - std::vector> expectedOutputId{{4}, {0}, {2}, {0}}; - std::vector> expectedOutputIds; - expectedOutputIds.reserve(expectedOutputId.size() * params.batchSize); - - for (auto const& id : expectedOutputId) - { - for (int32_t i = 0; i < params.batchSize; ++i) - { - expectedOutputIds.emplace_back(id); - } - } - this->runTest(expectedOutputIds, params); -} - -} // namespace diff --git a/cpp/tests/unit_tests/runtime/CMakeLists.txt b/cpp/tests/unit_tests/runtime/CMakeLists.txt index 3a171ee39877..9dd44ecbc110 100644 --- a/cpp/tests/unit_tests/runtime/CMakeLists.txt +++ b/cpp/tests/unit_tests/runtime/CMakeLists.txt @@ -15,10 +15,7 @@ add_gtest(bufferManagerTest bufferManagerTest.cpp) add_gtest(cudaMemPoolTest cudaMemPoolTest.cpp) -add_gtest(decodingLayerWorkspaceTest decodingLayerWorkspaceTest.cpp) add_gtest(gdrcopyTest gdrcopyTest.cpp) -add_gtest(gptDecoderBatchedTest gptDecoderBatchedTest.cpp) -add_gtest(gptDecoderTest gptDecoderTest.cpp) add_gtest(hostAccessibleDeviceAllocatorTest hostAccessibleDeviceAllocatorTest.cu) add_gtest(iBufferTest iBufferTest.cpp) @@ -32,7 +29,6 @@ add_gtest(moeLoadBalancerTest moeLoadBalancerTest.cpp) add_gtest(runtimeMpiUtilsTest mpiUtilsTest.cpp) add_gtest(runtimeKernelTest runtimeKernelTest.cpp) add_gtest(samplingConfigTest samplingConfigTest.cpp) -add_gtest(samplingTest samplingTest.cpp) add_gtest(sanitizerTest sanitizerTest.cpp) add_gtest(tllmBuffersTest tllmBuffersTest.cpp) add_gtest(transposeKVKernelTest transposeKVKernelTest.cpp) diff --git a/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp b/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp deleted file mode 100644 index 74f8faa37c87..000000000000 --- a/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp +++ /dev/null @@ -1,202 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/decodingLayerWorkspace.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/common/workspace.h" -#include -#include - -using namespace tensorrt_llm; - -namespace -{ -void populateCpuBufferWithRandomBytes(uint64_t seed, runtime::IBuffer& buffer) -{ - std::mt19937 generator(seed); - std::uniform_int_distribution distribution(0, 255); - auto* bufferPtr = reinterpret_cast(buffer.data()); - for (size_t i = 0; i < buffer.getSizeInBytes(); ++i) - { - *(bufferPtr + i) = static_cast(distribution(generator)); - } -} - -bool areMemoryRegionsEqual(void const* ptr1, void const* ptr2, size_t size) -{ - // Use std::memcmp to compare the memory regions - return std::memcmp(ptr1, ptr2, size) == 0; -} - -void testBufferEqual(runtime::IBuffer& left, runtime::IBuffer& right) -{ - auto const size = left.getSizeInBytes(); - ASSERT_EQ(size, right.getSizeInBytes()); - ASSERT_TRUE(areMemoryRegionsEqual(left.data(), right.data(), size)); -} -} // namespace - -auto const maxBatchSizePowersOfTwo = testing::Range(0, 14, 1); -auto const workspaceSizePowersOfTwo = testing::Range(0, 30, 2); - -auto const initialBatchAndWorkspaceSizes = testing::Combine(maxBatchSizePowersOfTwo, workspaceSizePowersOfTwo); - -using BasicUsageParamType = std::tuple; - -class BasicUsageTest : public testing::TestWithParam -{ - void SetUp() override - { - auto const deviceCount = common::getDeviceCount(); - if (deviceCount > 0) - { - mBufferManager = std::make_shared(std::make_unique()); - } - else - { - GTEST_SKIP() << "This test suite cannot run on systems with no devices."; - } - } - -protected: - std::shared_ptr mBufferManager = nullptr; -}; - -TEST_P(BasicUsageTest, TestBasicUsageOfDecodingLayerWorkspace) -{ - auto const [maxBatchSizePowerOfTwo, workspaceSizePowerOfTwo] = GetParam(); - auto const maxBatchSize = static_cast(std::pow(2, maxBatchSizePowerOfTwo)); - auto const workspaceSizeInBytes = static_cast(std::pow(2, workspaceSizePowerOfTwo)); - auto const decoderDomain = tensorrt_llm::layers::DecoderDomain(maxBatchSize, 1, 1000, 1024); - - // Testing constructing the workspace. - auto workspace = runtime::DecodingLayerWorkspace( - mBufferManager, decoderDomain, tensorrt_llm::runtime::TRTDataType::value, workspaceSizeInBytes); - mBufferManager->getStream().synchronize(); - ASSERT_EQ(workspace.getWorkspaceDeviceBuffer()->getSizeInBytes(), workspaceSizeInBytes) - << "The workspace size is not equal to the size we asked it to be."; - ASSERT_EQ(workspace.getDeviceBatchSlots()->getSize(), maxBatchSize) - << "The size of the device batch slots is not the max batch size provided to the workspace"; - - // Testing enlarging the workspace. - workspace.resize(workspaceSizeInBytes / 2); - ASSERT_EQ(workspace.getWorkspaceDeviceBuffer()->getSizeInBytes(), workspaceSizeInBytes) - << "The workspace size should not shrink."; - auto const biggerWorkspaceSize = workspaceSizeInBytes * 2; - workspace.resize(biggerWorkspaceSize); - ASSERT_EQ(workspace.getWorkspaceDeviceBuffer()->getSizeInBytes(), biggerWorkspaceSize) - << "The workspace was not enlarged as expected"; - - // Checking that the device batch slots are actually on device - auto const deviceBatchSlots = workspace.getDeviceBatchSlots(); - ASSERT_EQ(deviceBatchSlots->getMemoryType(), runtime::MemoryType::kGPU) - << "The device batch slots should be on device."; - - auto const* deviceBatchSlotsPtr = workspace.getDeviceBatchSlotsPtr(); - ASSERT_EQ(tensorrt_llm::common::getPtrCudaMemoryType(deviceBatchSlotsPtr), cudaMemoryType::cudaMemoryTypeDevice) - << "Pointer to device batch slots should have cudaMemoryType = device."; -} - -INSTANTIATE_TEST_SUITE_P(BasicUsage, BasicUsageTest, initialBatchAndWorkspaceSizes); - -auto const randomSeeds = testing::Values(static_cast(1234)); -auto const tensorDimensions = testing::Values(10, 100); -auto const tensorDataTypes = testing::Values(runtime::TRTDataType::value, runtime::TRTDataType::value, - runtime::TRTDataType::value, runtime::TRTDataType::value); -auto const tensorDataTypesTuples = testing::Combine(tensorDataTypes, tensorDataTypes, tensorDataTypes); - -auto const tensorShapeTuples = testing::Combine(tensorDimensions, tensorDimensions, tensorDimensions); -auto const mirrorInWorkspaceParams = testing::Combine(tensorDataTypesTuples, tensorShapeTuples, randomSeeds); - -using MirrorInWorkspaceParamType - = std::tuple, - std::tuple, std::uint64_t>; - -class MirrorInWorkspaceTest : public testing::TestWithParam -{ - void SetUp() override - { - auto const deviceCount = common::getDeviceCount(); - if (deviceCount > 0) - { - mBufferManager = std::make_shared(std::make_unique()); - } - else - { - GTEST_SKIP() << "This test suite cannot run on systems with no devices."; - } - } - -protected: - std::shared_ptr mBufferManager = nullptr; -}; - -TEST_P(MirrorInWorkspaceTest, TestMirrorInWorkspaceFunctionality) -{ - auto const [tensorDataTypes, tensorDimensions, randomSeed] = GetParam(); - auto const [tensorDataType1, tensorDataType2, tensorDataType3] = tensorDataTypes; - auto const [tensorDimension1, tensorDimension2, tensorDimension3] = tensorDimensions; - auto const decoderDomain = tensorrt_llm::layers::DecoderDomain(128, 1, 1000, 1024); - - // Testing constructing the workspace. - auto const hostTensorShape1 - = tensorrt_llm::runtime::ITensor::makeShape({tensorDimension1, tensorDimension2, tensorDimension3}); - auto const hostTensorShape2 - = tensorrt_llm::runtime::ITensor::makeShape({tensorDimension2, tensorDimension3, tensorDimension1}); - auto const hostTensorShape3 - = tensorrt_llm::runtime::ITensor::makeShape({tensorDimension3, tensorDimension1, tensorDimension2}); - runtime::ITensor::SharedPtr const hostTensor1 = mBufferManager->cpu(hostTensorShape1, tensorDataType1); - runtime::ITensor::SharedPtr const hostTensor2 = mBufferManager->cpu(hostTensorShape1, tensorDataType2); - runtime::ITensor::SharedPtr const hostTensor3 = mBufferManager->cpu(hostTensorShape1, tensorDataType3); - - auto const requiredWorkspaceSize = tensorrt_llm::runtime::DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(hostTensorShape1, tensorDataType1), std::make_pair(hostTensorShape2, tensorDataType2), - std::make_pair(hostTensorShape3, tensorDataType3)); - auto workspace = runtime::DecodingLayerWorkspace( - mBufferManager, decoderDomain, tensorrt_llm::runtime::TRTDataType::value, requiredWorkspaceSize); - mBufferManager->getStream().synchronize(); - - ASSERT_LE(hostTensor1->getSizeInBytes() + hostTensor2->getSizeInBytes() + hostTensor3->getSizeInBytes(), - requiredWorkspaceSize) - << "The calculated workspace size cannot possibly be enough to contain all the tensors."; - - constexpr std::size_t addressAlignment = tensorrt_llm::common::kCudaMemAlign; - constexpr std::size_t numTensors = 3; - constexpr std::size_t maxAlignmentOverhead = numTensors * addressAlignment; - ASSERT_GE(hostTensor1->getSizeInBytes() + hostTensor2->getSizeInBytes() + hostTensor3->getSizeInBytes() - + maxAlignmentOverhead, - requiredWorkspaceSize) - << "We probably overestimate the amount of space the workspace requires."; - populateCpuBufferWithRandomBytes(randomSeed, *hostTensor1); - populateCpuBufferWithRandomBytes(randomSeed, *hostTensor2); - populateCpuBufferWithRandomBytes(randomSeed, *hostTensor3); - - auto const [deviceTensor1, deviceTensor2, deviceTensor3] - = workspace.mirrorInWorkspace(hostTensor1, hostTensor2, hostTensor3); - - runtime::ITensor::SharedPtr const hostTensorCopy1 = mBufferManager->cpu(hostTensorShape1, tensorDataType1); - runtime::ITensor::SharedPtr const hostTensorCopy2 = mBufferManager->cpu(hostTensorShape1, tensorDataType2); - runtime::ITensor::SharedPtr const hostTensorCopy3 = mBufferManager->cpu(hostTensorShape1, tensorDataType3); - mBufferManager->copy(*deviceTensor1, *hostTensorCopy1); - mBufferManager->copy(*deviceTensor2, *hostTensorCopy2); - mBufferManager->copy(*deviceTensor3, *hostTensorCopy3); - testBufferEqual(*hostTensor1, *hostTensorCopy1); - testBufferEqual(*hostTensor2, *hostTensorCopy2); - testBufferEqual(*hostTensor3, *hostTensorCopy3); -} - -INSTANTIATE_TEST_SUITE_P(MirrorInWorkspace, MirrorInWorkspaceTest, mirrorInWorkspaceParams); diff --git a/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp b/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp deleted file mode 100644 index a979c9a4699f..000000000000 --- a/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp +++ /dev/null @@ -1,837 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/runtime/gptDecoderBatched.h" -#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include - -#include -#include -#include - -using namespace tensorrt_llm::runtime; - -namespace tle = tensorrt_llm::executor; -namespace tc = tensorrt_llm::common; -namespace tb = tensorrt_llm::batch_manager; - -using TensorPtr = ITensor::SharedPtr; - -namespace -{ - -// Local copy of the former MakeDecodingBatchInputOutput::createDecoderBatchInputs -// helper, which was removed with the TensorRT-engine execution path. The decoder -// under test is backend-agnostic; this builds its step-batched inputs directly. -void createDecoderBatchInputs(tb::DecoderInputBuffers& inputBuffers, std::vector const& activeSlots, - decoder::DecoderState const& decoderState) -{ - auto const& numDecodingEngineTokens = decoderState.getNumDecodingEngineTokens(); - auto const& maxDecodingEngineTokens = decoderState.getMaxDecodingEngineTokens(); - auto const& maxDecodingDecoderTokens = decoderState.getMaxDecodingDecoderTokens(); - auto const maxDecoderSteps = tc::ceilDiv(maxDecodingEngineTokens, maxDecodingDecoderTokens); - - auto& batchSlots = inputBuffers.forwardBatchSlots; - auto& decoderLogits = inputBuffers.decoderLogits; - - for (SizeType32 step = 0; step < maxDecoderSteps; ++step) - { - batchSlots.at(step)->resize(activeSlots.size()); - } - - auto constexpr singleRequest = 1; - - std::vector batchSizes(maxDecoderSteps); - std::vector> batchLogits(maxDecoderSteps); - auto maxActiveDecoderSteps = 1; - for (size_t batchIdx = 0; batchIdx < activeSlots.size(); ++batchIdx) - { - auto const slot = activeSlots.at(batchIdx); - auto const& logits = decoderLogits.at(batchIdx); - - auto const numDecoderSteps = tc::ceilDiv(numDecodingEngineTokens.at(slot), maxDecodingDecoderTokens); - maxActiveDecoderSteps = std::max(maxActiveDecoderSteps, numDecoderSteps); - for (SizeType32 step = 0; step < numDecoderSteps; ++step) - { - auto batchSlotsRange = BufferRange(*batchSlots.at(step)); - batchSlotsRange[batchSizes[step]] = slot; - batchSizes[step]++; - auto logitsSlice = ITensor::slice(logits, step, singleRequest); - batchLogits[step].emplace_back(std::move(logitsSlice)); - } - } - - for (SizeType32 step = 0; step < maxDecoderSteps; ++step) - { - batchSlots.at(step)->resize(batchSizes[step]); - } - batchLogits.resize(maxActiveDecoderSteps); - - inputBuffers.maxDecoderSteps = maxActiveDecoderSteps; - inputBuffers.batchLogits = batchLogits; -} - -} // namespace - -namespace -{ - -std::shared_ptr createLlmRequest(SizeType32 batchSlot, SizeType32 inputLengths, - SizeType32 generatedTokensPerSteps, SizeType32 acceptedTokensPerStep, TokenIdType inputTokenId, - TokenIdType expectedTokenId, SizeType32 maxNewTokens, SamplingConfig const& samplingConfig, TokenIdType endId) -{ - auto constexpr requestId = 0; - auto inputTokens = std::make_shared(inputLengths, inputTokenId); - bool isStreaming = false; - auto request - = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, endId); - request->mSeqSlot = batchSlot; - - if (generatedTokensPerSteps > 1) - { - TokenIdType constexpr tokenToReject{1}; - TLLM_CHECK(tokenToReject != expectedTokenId); - // fill with tokens to reject - auto draftTokens = std::make_shared(generatedTokensPerSteps - 1, tokenToReject); - std::fill(draftTokens->begin(), draftTokens->begin() + acceptedTokensPerStep, expectedTokenId); - request->setDraftTokens(draftTokens); - } - - return request; -} - -std::vector> createLlmRequests(std::vector const& inputLengths, - std::vector const& generatedTokensPerSteps, std::vector const& acceptedTokensPerStep, - TokenIdType inputTokenId, TokenIdType expectedTokenId, TensorPtr const& batchSlots, - std::vector const& allSamplingConfigs, SizeType32 maxNewTokens, SizeType32 endId) -{ - auto batchSlotsRange = BufferRange(*batchSlots); - auto const localBatchSize = batchSlots->getSize(); - - std::vector> requests; - for (size_t bi = 0; bi < localBatchSize; ++bi) - { - auto const batchSlot = batchSlotsRange[bi]; - auto llmReq = createLlmRequest(batchSlot, inputLengths[batchSlot], generatedTokensPerSteps[batchSlot], - acceptedTokensPerStep[batchSlot], inputTokenId, expectedTokenId, maxNewTokens, - allSamplingConfigs[batchSlot], endId); - requests.emplace_back(std::move(llmReq)); - } - - return requests; -} - -void newRequests(std::vector> const& requests, TensorPtr const& batchSlots, - tensorrt_llm::DataType logitsType, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - tle::DecodingConfig const& decodingConfig, GptDecoderBatched& decoder, CudaStream const& runtimeStream, - SizeType32 maxSequenceLength, tb::DecoderInputBuffers& inputBuffers, decoder::DecoderState& decoderState) -{ - auto const& decoderStream = *decoder.getDecoderStream(); - - 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, - logitsType, modelConfig, worldConfig, runtimeStream, decoderStream, maxSequenceLength, std::nullopt); - - std::vector samplingConfigs; - samplingConfigs.reserve(requests.size()); - for (auto const& llmReq : requests) - { - samplingConfigs.emplace_back(llmReq->mSamplingConfig); - } - - // Setup underlying decoder. - auto samplingConfig = SamplingConfig(samplingConfigs); - decoder.getUnderlyingDecoder().setup( - samplingConfig, localBatchSize, batchSlots, {decoderState.getJointDecodingOutput()}); - - CudaEvent event{}; - decoderStream.record(event); - runtimeStream.wait(event); -} - -void createDecoderInputs(tb::DecoderInputBuffers& inputBuffers, SizeType32 batchSize, SizeType32 vocabSizePadded, - tensorrt_llm::DataType dataType, std::vector& samplingConfigs, - std::vector const& generatedTokensPerSteps, bool computeLogProbs, BufferManager& manager) -{ - auto& logits = inputBuffers.decoderLogits; - logits.reserve(batchSize); - for (auto batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - auto& samplingConfig = samplingConfigs[batchIdx]; - auto const beamWidth = samplingConfig.beamWidth; - samplingConfig.outputLogProbs = {{computeLogProbs}}; - samplingConfig.cumLogProbs = {{computeLogProbs}}; - - logits.emplace_back( - manager.gpu(ITensor::makeShape({generatedTokensPerSteps[batchIdx], beamWidth, vocabSizePadded}), dataType)); - manager.setZero(*logits.back()); - } -} - -void copySequenceLengths( - std::vector const& tiledInputLengths, ITensor& sequenceLengths, BufferManager const& manager) -{ - TLLM_CHECK(sequenceLengths.getSize() == tiledInputLengths.size()); - manager.copy(tiledInputLengths.data(), sequenceLengths); -} - -[[nodiscard]] std::vector getFinished( - ITensor const& finishedSum, std::vector const& samplingConfigs, BufferManager& manager) -{ - auto finishedSumHost = manager.copyFrom(finishedSum, MemoryType::kCPU); - auto finishedSumHostRange = BufferRange(*finishedSumHost); - std::vector finished(finishedSumHostRange.size()); - std::transform(finishedSumHostRange.begin(), finishedSumHostRange.end(), samplingConfigs.begin(), finished.begin(), - [](SizeType32 sum, SamplingConfig const& config) { return sum == config.beamWidth; }); - - return finished; -} - -void advanceSequenceLengths(std::vector& sequenceLengths, - std::vector const& acceptedTokensPerStep, std::vector const& samplingConfigs, - std::vector const& finished, SizeType32 batchSize, SizeType32 maxBeamWidth) -{ - for (int batchIdx = 0; batchIdx < batchSize; batchIdx++) - { - if (!finished.at(batchIdx)) - { - for (int beamId = 0; beamId < samplingConfigs.at(batchIdx).beamWidth; beamId++) - { - sequenceLengths.at(tc::flat_index2(batchIdx, beamId, maxBeamWidth)) - += acceptedTokensPerStep.at(batchIdx) + 1; - } - } - } -} - -void checkSequenceLengths( - ITensor const& sequenceLengths, std::vector const& expectedLengths, BufferManager& manager) -{ - auto sequenceLengthsHost = manager.copyFrom(sequenceLengths, MemoryType::kCPU); - auto sequenceLengthsHostRange = BufferRange(*sequenceLengthsHost); - EXPECT_THAT(sequenceLengthsHostRange, ::testing::ElementsAreArray(expectedLengths)); -} - -void verifyResults(BufferManager& manager, decoder::DecoderState const& decoderState, - std::vector const& samplingConfigs, std::vector const& inputLengths, - std::vector const& sequenceLengths, SizeType32 batchSize, SizeType32 maxBeamWidth, - SizeType32 maxSeqLength, SizeType32 inputTokenId, SizeType32 expectedTokenId, SizeType32 endId) -{ - for (auto b = 0; b < batchSize; ++b) - { - auto outputsIds = decoderState.getIds(b); - // TODO: test parentIds - // parentIds = decoder.getParentIds(); - ASSERT_TRUE(outputsIds); - auto outputShape = outputsIds->getShape(); - EXPECT_EQ(outputShape.nbDims, 2); - EXPECT_EQ(outputShape.d[0], maxBeamWidth); - EXPECT_EQ(outputShape.d[1], maxSeqLength); - - auto outputsIdsHost = manager.copyFrom(*outputsIds, MemoryType::kCPU); - auto* output = bufferCast(*outputsIdsHost); - - auto const& samplingConfig = samplingConfigs.at(b); - - for (auto beamIndex = 0; beamIndex < samplingConfig.beamWidth; ++beamIndex) - { - auto const result = (samplingConfig.beamWidth == 1) ? expectedTokenId : beamIndex; - - auto* const outputPtr = output + tc::flat_index(outputShape.d, beamIndex, 0); - - auto const inputLength = inputLengths.at(b); - auto* begin = outputPtr; - auto* end = outputPtr + inputLength; - ASSERT_LE(begin, end) << "bad input length " << inputLength; - - // Only the first beam contains the input token ID, the other beams point to it. - if (beamIndex == 0) - { - ASSERT_THAT(std::vector(begin, end), ::testing::Each(inputTokenId)) - << "input tokens: " - << "b:" << b << " bw: " << beamIndex; - } - - auto const seqLength = sequenceLengths.at(tc::flat_index2(b, beamIndex, maxBeamWidth)); - begin = end; - end = outputPtr + seqLength; - ASSERT_LE(begin, end) << "bad seq length " << seqLength; - ASSERT_THAT(std::vector(begin, end), ::testing::Each(result)) << "new tokens: " - << "b:" << b << " bw: " << beamIndex; - begin = end; - end = outputPtr + maxSeqLength; - ASSERT_LE(begin, end) << "bad max length " << maxSeqLength; - ASSERT_THAT(std::vector(begin, end), ::testing::Each(endId)) << "padding: " - << "b:" << b << " bw: " << beamIndex; - } - } -} - -void testDecoder(tensorrt_llm::DataType const dtype, std::vector& samplingConfigs, - SizeType32 maxBeamWidth, bool computeLogProbs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - SizeType32 constexpr tensorParallelism{1}; - SizeType32 constexpr pipelineParallelism{1}; - SizeType32 constexpr contextParallelism{1}; - SizeType32 constexpr localRank{0}; - WorldConfig const worldConfig{tensorParallelism, pipelineParallelism, contextParallelism, localRank}; - - SizeType32 constexpr vocabSize{51200}; - SizeType32 constexpr nbAttentionLayers{2}; - SizeType32 constexpr nbRnnLayers{0}; - SizeType32 constexpr nbHeads{16}; - SizeType32 constexpr hiddenSize{1024}; - ModelConfig modelConfig{ - vocabSize, nbAttentionLayers + nbRnnLayers, nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, dtype}; - modelConfig.useGptAttentionPlugin(false); - - auto streamPtr = std::make_shared(); - BufferManager manager(streamPtr); - - TokenIdType constexpr endId{50257}; - - auto const dataType = modelConfig.getDataType(); - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const batchSize = static_cast(samplingConfigs.size()); - SizeType32 constexpr maxInputLength{8}; - SizeType32 const maxNewTokens{2}; - auto const maxSeqLength = maxInputLength + maxNewTokens; - SizeType32 constexpr maxGeneratedTokensPerStep{1}; - - std::vector inputLengths(batchSize); - std::iota(inputLengths.begin(), inputLengths.end(), 4); - - std::vector tiledInputLengths; - for (int batchIdx = 0; batchIdx < inputLengths.size(); batchIdx++) - { - for (int beamId = 0; beamId < maxBeamWidth; beamId++) - { - tiledInputLengths.push_back(inputLengths.at(batchIdx)); - } - } - - std::vector generatedTokensPerSteps(batchSize); - std::vector acceptedTokensPerStep(batchSize); - for (auto batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - generatedTokensPerSteps[batchIdx] = maxGeneratedTokensPerStep; - acceptedTokensPerStep[batchIdx] = generatedTokensPerSteps[batchIdx] - 1; - } - - auto constexpr inputTokenId = 1; - auto constexpr expectedTokenId = 1023; - - // We set maxAttentionWindow = maxSeqLength, but it can be smaller than maxSeqLength (cyclic kv cache). - auto const maxAttentionWindow = maxSeqLength; - SizeType32 const sinkTokenLength{0}; - - auto const decodingMode = maxBeamWidth == 1 ? tle::DecodingMode::TopKTopP() : tle::DecodingMode::BeamSearch(); - tle::DecodingConfig decodingConfig{decodingMode}; - - // set up decoder - auto decoder = GptDecoderBatched(streamPtr); - decoder.setup(decodingMode, batchSize, maxBeamWidth, dataType, modelConfig, worldConfig); - - decoder::DecoderState decoderState; - decoderState.setup(batchSize, maxBeamWidth, maxAttentionWindow, sinkTokenLength, maxSeqLength, dataType, - modelConfig, worldConfig, manager); - - // set up inputs and outputs - tb::DecoderInputBuffers inputBuffers(batchSize, maxGeneratedTokensPerStep, manager); - auto batchSlotsRange = BufferRange(*inputBuffers.setupBatchSlots); - std::iota(batchSlotsRange.begin(), batchSlotsRange.end(), 0); - - createDecoderInputs(inputBuffers, batchSize, vocabSizePadded, dataType, samplingConfigs, generatedTokensPerSteps, - computeLogProbs, manager); - manager.setZero(*decoderState.getCacheIndirectionInput()); - copySequenceLengths(tiledInputLengths, *decoderState.getSequenceLengths(), manager); - - auto requests = createLlmRequests(inputLengths, generatedTokensPerSteps, acceptedTokensPerStep, inputTokenId, - expectedTokenId, inputBuffers.setupBatchSlots, samplingConfigs, maxNewTokens, endId); - newRequests(requests, inputBuffers.setupBatchSlots, dataType, modelConfig, worldConfig, decodingConfig, decoder, - *streamPtr, maxSeqLength, inputBuffers, decoderState); - cudaDeviceSynchronize(); - - auto expectedLengths = tiledInputLengths; - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - - auto const& finished = getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager); - EXPECT_EQ(finished.size(), batchSize); - EXPECT_THAT(finished, ::testing::Each(false)); - - verifyResults(manager, decoderState, samplingConfigs, inputLengths, expectedLengths, batchSize, maxBeamWidth, - maxSeqLength, inputTokenId, expectedTokenId, endId); - - // run decoder for 1 step - advanceSequenceLengths(expectedLengths, acceptedTokensPerStep, samplingConfigs, - getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), batchSize, maxBeamWidth); - - auto activeSlots = std::vector(batchSize); - std::iota(activeSlots.begin(), activeSlots.end(), 0); - createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); - decoder.forward(decoderState, inputBuffers); - - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - EXPECT_THAT(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), ::testing::Each(false)); - - verifyResults(manager, decoderState, samplingConfigs, inputLengths, expectedLengths, batchSize, maxBeamWidth, - maxSeqLength, inputTokenId, expectedTokenId, endId); - - // run decoder for 1 step - advanceSequenceLengths(expectedLengths, acceptedTokensPerStep, samplingConfigs, - getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), batchSize, maxBeamWidth); - decoder.forward(decoderState, inputBuffers); - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - EXPECT_THAT(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), ::testing::Each(true)); - - verifyResults(manager, decoderState, samplingConfigs, inputLengths, expectedLengths, batchSize, maxBeamWidth, - maxSeqLength, inputTokenId, expectedTokenId, endId); - - EXPECT_NO_THROW(decoder.forward(decoderState, inputBuffers)); - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - - TensorPtr batchSlotsView = ITensor::slice(inputBuffers.setupBatchSlots, 0, 1); - requests = createLlmRequests(inputLengths, generatedTokensPerSteps, acceptedTokensPerStep, inputTokenId, - expectedTokenId, batchSlotsView, samplingConfigs, maxNewTokens, endId); - newRequests(requests, batchSlotsView, dataType, modelConfig, worldConfig, decodingConfig, decoder, *streamPtr, - maxSeqLength, inputBuffers, decoderState); - EXPECT_FALSE(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager)[0]); -} - -void testDecoderWavefront(tensorrt_llm::DataType const dtype, std::vector& samplingConfigs, - SizeType32 maxBeamWidth, bool computeLogProbs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - SizeType32 constexpr tensorParallelism{1}; - SizeType32 constexpr pipelineParallelism{1}; - SizeType32 constexpr contextParallelism{1}; - SizeType32 constexpr localRank{0}; - WorldConfig const worldConfig{tensorParallelism, pipelineParallelism, contextParallelism, localRank}; - - SizeType32 constexpr vocabSize{51200}; - SizeType32 constexpr nbAttentionLayers{2}; - SizeType32 constexpr nbRnnLayers{0}; - SizeType32 constexpr nbHeads{16}; - SizeType32 constexpr hiddenSize{1024}; - ModelConfig modelConfig{ - vocabSize, nbAttentionLayers + nbRnnLayers, nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, dtype}; - modelConfig.useGptAttentionPlugin(false); - - auto streamPtr = std::make_shared(); - BufferManager manager(streamPtr); - - TokenIdType constexpr endId{50257}; - - auto const dataType = modelConfig.getDataType(); - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const batchSize = static_cast(samplingConfigs.size()); - SizeType32 constexpr maxInputLength{8}; - SizeType32 constexpr maxNewTokens{8}; - auto constexpr maxSeqLength = maxInputLength + maxNewTokens; - SizeType32 constexpr maxGeneratedTokensPerStep{1}; - - std::vector inputLengths(batchSize); - std::iota(inputLengths.begin(), inputLengths.end(), 4); - - std::vector tiledInputLengths; - for (SizeType32 const inputLength : inputLengths) - { - for (SizeType32 beamId = 0; beamId < maxBeamWidth; beamId++) - { - tiledInputLengths.push_back(inputLength); - } - } - - std::vector generatedTokensPerSteps(batchSize); - std::vector acceptedTokensPerStep(batchSize); - for (auto batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - generatedTokensPerSteps[batchIdx] = maxGeneratedTokensPerStep; - acceptedTokensPerStep[batchIdx] = generatedTokensPerSteps[batchIdx] - 1; - } - - auto constexpr inputTokenId = 1; - auto constexpr expectedTokenId = 1023; - - // We set maxAttentionWindow = maxSeqLength, but it can be smaller than maxSeqLength (cyclic kv cache). - auto const maxAttentionWindow = maxSeqLength; - SizeType32 const sinkTokenLength{0}; - - auto const decodingMode = maxBeamWidth == 1 ? tle::DecodingMode::TopKTopP() : tle::DecodingMode::BeamSearch(); - tle::DecodingConfig decodingConfig{decodingMode}; - - // set up decoder - auto decoder = GptDecoderBatched(streamPtr); - decoder.setup(decodingMode, batchSize, maxBeamWidth, dataType, modelConfig, worldConfig); - - decoder::DecoderState decoderState; - decoderState.setup(batchSize, maxBeamWidth, maxAttentionWindow, sinkTokenLength, maxSeqLength, dataType, - modelConfig, worldConfig, manager); - - // set up inputs and outputs - tb::DecoderInputBuffers inputBuffers(batchSize, maxGeneratedTokensPerStep, manager); - - createDecoderInputs(inputBuffers, batchSize, vocabSizePadded, dataType, samplingConfigs, generatedTokensPerSteps, - computeLogProbs, manager); - manager.setZero(*decoderState.getCacheIndirectionInput()); - copySequenceLengths(tiledInputLengths, *decoderState.getSequenceLengths(), manager); - - std::vector const expectedSteps(batchSize, 0); - auto expectedLengths = tiledInputLengths; - - auto const& finished = getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager); - EXPECT_EQ(finished.size(), batchSize); - std::vector expectedFinished(batchSize, false); - - auto batchSlotsRange = BufferRange(*inputBuffers.setupBatchSlots); - std::iota(batchSlotsRange.begin(), batchSlotsRange.end(), 0); - - for (auto batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - TensorPtr const newBatchSlot = ITensor::slice(inputBuffers.setupBatchSlots, batchIdx, 1); - auto requests = createLlmRequests(inputLengths, generatedTokensPerSteps, acceptedTokensPerStep, inputTokenId, - expectedTokenId, newBatchSlot, samplingConfigs, maxNewTokens, endId); - newRequests(requests, newBatchSlot, dataType, modelConfig, worldConfig, decodingConfig, decoder, *streamPtr, - maxSeqLength, inputBuffers, decoderState); - - auto activeSlots = std::vector(batchIdx + 1); - std::iota(activeSlots.begin(), activeSlots.end(), 0); - createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); - decoder.forward(decoderState, inputBuffers); - - advanceSequenceLengths( - expectedLengths, acceptedTokensPerStep, samplingConfigs, expectedFinished, batchIdx + 1, maxBeamWidth); - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - - for (auto bi = 0; bi <= batchIdx; ++bi) - { - auto firstBeamIndex = tc::flat_index2(bi, 0, maxBeamWidth); - expectedFinished.at(bi) - = expectedLengths.at(firstBeamIndex) - tiledInputLengths.at(firstBeamIndex) >= maxNewTokens; - } - EXPECT_THAT(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), - ::testing::ElementsAreArray(expectedFinished)); - } - - auto activeSlots = std::vector(batchSize); - std::iota(activeSlots.begin(), activeSlots.end(), 0); - auto finishedVec = getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager); - while (!std::all_of(expectedFinished.begin(), expectedFinished.end(), [](bool finish) { return finish; })) - { - createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); - decoder.forward(decoderState, inputBuffers); - finishedVec = getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager); - - advanceSequenceLengths( - expectedLengths, acceptedTokensPerStep, samplingConfigs, expectedFinished, batchSize, maxBeamWidth); - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - - for (auto bi = 0; bi < batchSize; ++bi) - { - auto firstBeamIndex = tc::flat_index2(bi, 0, maxBeamWidth); - expectedFinished.at(bi) - = expectedLengths.at(firstBeamIndex) - tiledInputLengths.at(firstBeamIndex) >= maxNewTokens; - } - EXPECT_THAT(finishedVec, ::testing::ElementsAreArray(expectedFinished)); - - activeSlots.clear(); - for (auto batchIdx = 0; batchIdx < batchSize; ++batchIdx) - { - if (!finishedVec.at(batchIdx)) - { - activeSlots.push_back(batchIdx); - } - } - } - - verifyResults(manager, decoderState, samplingConfigs, inputLengths, expectedLengths, batchSize, maxBeamWidth, - maxSeqLength, inputTokenId, expectedTokenId, endId); -} - -void testDecoderDraft(tensorrt_llm::DataType const dtype, std::vector& samplingConfigs, - SizeType32 maxBeamWidth, std::vector const& generatedTokensPerSteps, - std::vector const& acceptedTokensPerStep, SizeType32 maxGeneratedTokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(maxBeamWidth == 1); - - SizeType32 constexpr tensorParallelism{1}; - SizeType32 constexpr pipelineParallelism{1}; - SizeType32 constexpr contextParallelism{1}; - SizeType32 constexpr localRank{0}; - WorldConfig const worldConfig{tensorParallelism, pipelineParallelism, contextParallelism, localRank}; - - SizeType32 constexpr vocabSize{51200}; - SizeType32 constexpr nbAttentionLayers{2}; - SizeType32 constexpr nbRnnLayers{0}; - SizeType32 constexpr nbHeads{16}; - SizeType32 constexpr hiddenSize{1024}; - ModelConfig modelConfig{ - vocabSize, nbAttentionLayers + nbRnnLayers, nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, dtype}; - modelConfig.useGptAttentionPlugin(false); - modelConfig.setSpeculativeDecodingMode(SpeculativeDecodingMode::DraftTokensExternal()); - - auto streamPtr = std::make_shared(); - BufferManager manager(streamPtr); - - TokenIdType constexpr endId{50257}; - - auto const dataType = modelConfig.getDataType(); - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - auto const batchSize = static_cast(samplingConfigs.size()); - SizeType32 constexpr maxInputLength{8}; - SizeType32 const maxNewTokens{4}; - auto const maxSeqLength = maxInputLength + maxNewTokens; - - std::vector inputLengths(batchSize); - std::iota(inputLengths.begin(), inputLengths.end(), 4); - - std::vector tiledInputLengths; - for (int batchIdx = 0; batchIdx < inputLengths.size(); batchIdx++) - { - for (int beamId = 0; beamId < maxBeamWidth; beamId++) - { - tiledInputLengths.push_back(inputLengths.at(batchIdx)); - } - } - - auto constexpr inputTokenId = 1; - auto constexpr expectedTokenId = 1023; - - // We set maxAttentionWindow = maxSeqLength, but it can be smaller than maxSeqLength (cyclic kv cache). - auto const maxAttentionWindow = maxSeqLength; - SizeType32 const sinkTokenLength{0}; - - auto const decodingMode = tle::DecodingMode::ExternalDraftTokens(); // only supports bw=1 - tle::DecodingConfig decodingConfig{decodingMode}; - - // set up decoder - auto decoder = GptDecoderBatched(streamPtr); - decoder.setup(decodingMode, batchSize, maxBeamWidth, dataType, modelConfig, worldConfig); - - decoder::DecoderState decoderState; - decoderState.setup(batchSize, maxBeamWidth, maxAttentionWindow, sinkTokenLength, maxSeqLength, dataType, - modelConfig, worldConfig, manager); - if (!modelConfig.getSpeculativeDecodingMode().isNone()) - { - decoderState.setupSpeculativeDecoding(modelConfig.getSpeculativeDecodingMode(), maxGeneratedTokensPerStep, - dtype, modelConfig, worldConfig, manager); - } - - // set up inputs and outputs - tb::DecoderInputBuffers inputBuffers(batchSize, maxGeneratedTokensPerStep, manager); - - createDecoderInputs( - inputBuffers, batchSize, vocabSizePadded, dataType, samplingConfigs, generatedTokensPerSteps, false, manager); - manager.setZero(*decoderState.getCacheIndirectionInput()); - copySequenceLengths(tiledInputLengths, *decoderState.getSequenceLengths(), manager); - - auto batchSlotsRange = BufferRange(*inputBuffers.setupBatchSlots); - std::iota(batchSlotsRange.begin(), batchSlotsRange.end(), 0); - - auto requests = createLlmRequests(inputLengths, generatedTokensPerSteps, acceptedTokensPerStep, inputTokenId, - expectedTokenId, inputBuffers.setupBatchSlots, samplingConfigs, maxNewTokens, endId); - newRequests(requests, inputBuffers.setupBatchSlots, dataType, modelConfig, worldConfig, decodingConfig, decoder, - *streamPtr, maxSeqLength, inputBuffers, decoderState); - cudaDeviceSynchronize(); - - auto expectedLengths = tiledInputLengths; - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - - auto const& finished = getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager); - EXPECT_EQ(finished.size(), batchSize); - EXPECT_THAT(finished, ::testing::Each(false)); - - verifyResults(manager, decoderState, samplingConfigs, inputLengths, expectedLengths, batchSize, maxBeamWidth, - maxSeqLength, inputTokenId, expectedTokenId, endId); - - // run decoder for 1 step - advanceSequenceLengths(expectedLengths, acceptedTokensPerStep, samplingConfigs, - getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), batchSize, maxBeamWidth); - - auto activeSlots = std::vector(batchSize); - std::iota(activeSlots.begin(), activeSlots.end(), 0); - createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); - decoder.forward(decoderState, inputBuffers); - checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); - EXPECT_THAT(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), ::testing::Each(false)); - - verifyResults(manager, decoderState, samplingConfigs, inputLengths, expectedLengths, batchSize, maxBeamWidth, - maxSeqLength, inputTokenId, expectedTokenId, endId); -} - -} // namespace - -struct BeamConfig -{ - SizeType32 maxBeamWidth; - std::vector beamWidths; -}; - -using ParamType = std::tuple; - -std::string generateTestName(testing::TestParamInfo const& info) -{ - std::string name{std::get<0>(info.param) == tensorrt_llm::DataType::kFLOAT ? "Float" : "Half"}; - BeamConfig const beamConfig = std::get<1>(info.param); - name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); - for (auto const beamWdith : beamConfig.beamWidths) - { - name.append("Bw" + std::to_string(beamWdith)); - } - bool const computeLogProbs{std::get<2>(info.param)}; - if (computeLogProbs) - { - name.append("LogProbs"); - } - return name; -} - -class ParamTest : public ::testing::TestWithParam -{ -}; - -TEST_P(ParamTest, Test) -{ - tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; - BeamConfig const beamConfig{std::get<1>(GetParam())}; - bool const computeLogProbs{std::get<2>(GetParam())}; - std::vector samplingConfigs; - for (auto const beamWidth : beamConfig.beamWidths) - { - samplingConfigs.emplace_back(beamWidth); - } - - testDecoder(dtype, samplingConfigs, beamConfig.maxBeamWidth, computeLogProbs); -} - -INSTANTIATE_TEST_SUITE_P(DecoderBwTest, ParamTest, - testing::Combine(testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), - testing::Values(BeamConfig{1, {1, 1, 1}}, BeamConfig{3, {3, 3, 3, 3}}, BeamConfig{4, {4, 4, 4}}, - BeamConfig{10, {10, 10, 10}}), - testing::Values(false, true)), - generateTestName); - -class ParamWavefrontTest : public ::testing::TestWithParam -{ -}; - -TEST_P(ParamWavefrontTest, Test) -{ - tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; - BeamConfig const beamConfig{std::get<1>(GetParam())}; - bool const computeLogProbs{std::get<2>(GetParam())}; - bool const normalizeLogProbs{true}; - std::vector samplingConfigs; - for (auto const beamWidth : beamConfig.beamWidths) - { - samplingConfigs.emplace_back(beamWidth); - } - - testDecoderWavefront(dtype, samplingConfigs, beamConfig.maxBeamWidth, computeLogProbs); -} - -INSTANTIATE_TEST_SUITE_P(DecoderBwTest, ParamWavefrontTest, - testing::Combine(testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), - testing::Values(BeamConfig{1, {1, 1, 1}}, BeamConfig{3, {3, 3, 3, 3}}, BeamConfig{4, {4, 4, 4}}, - BeamConfig{10, {10, 10, 10}}), - testing::Values(false, true)), - generateTestName); - -struct DraftConfig -{ - SizeType32 maxGeneratedTokensPerStep; - std::vector generatedTokensPerSteps; - std::vector acceptedTokensPerStep; -}; - -using DraftTestParamType = std::tuple; - -class ParamDraftTest : public ::testing::TestWithParam -{ -}; - -TEST_P(ParamDraftTest, Test) -{ - tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; - BeamConfig const beamConfig{std::get<1>(GetParam())}; - DraftConfig const draftConfig{std::get<2>(GetParam())}; - - ASSERT_EQ(beamConfig.beamWidths.size(), draftConfig.acceptedTokensPerStep.size()); - ASSERT_EQ(beamConfig.beamWidths.size(), draftConfig.generatedTokensPerSteps.size()); - - std::vector samplingConfigs; - for (auto const beamWidth : beamConfig.beamWidths) - { - samplingConfigs.emplace_back(beamWidth); - } - - testDecoderDraft(dtype, samplingConfigs, beamConfig.maxBeamWidth, draftConfig.generatedTokensPerSteps, - draftConfig.acceptedTokensPerStep, draftConfig.maxGeneratedTokensPerStep); -} - -INSTANTIATE_TEST_SUITE_P(DecoderTest, ParamDraftTest, - testing::Combine(testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), - testing::Values(BeamConfig{1, {1, 1, 1}}), - testing::Values( // - DraftConfig{2, {1, 1, 1}, {0, 0, 0}}, DraftConfig{2, {2, 2, 2}, {1, 1, 1}}, - DraftConfig{4, {1, 2, 3}, {0, 0, 1}} - - )), - [](testing::TestParamInfo const& info) - { - std::string name{std::get<0>(info.param) == tensorrt_llm::DataType::kFLOAT ? "Float" : "Half"}; - BeamConfig const beamConfig = std::get<1>(info.param); - DraftConfig const draftConfig = std::get<2>(info.param); - name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); - auto const batchSize = beamConfig.beamWidths.size(); - for (auto const beamWdith : beamConfig.beamWidths) - { - name.append("Bw" + std::to_string(beamWdith)); - } - name.append("PerStep" + std::to_string(draftConfig.maxGeneratedTokensPerStep)); - for (std::size_t i = 0; i < batchSize; ++i) - { - auto const acc = draftConfig.acceptedTokensPerStep.at(i); - auto const gen = draftConfig.generatedTokensPerSteps.at(i); - name.append("Acc" + std::to_string(acc) + "of" + std::to_string(gen)); - } - return name; - }); diff --git a/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp b/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp deleted file mode 100644 index 5f620aa4d9c4..000000000000 --- a/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp +++ /dev/null @@ -1,271 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include - -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/gptDecoder.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -using namespace tensorrt_llm::runtime; - -namespace tc = tensorrt_llm::common; -namespace tle = tensorrt_llm::executor; - -namespace -{ - -bool forwardAndSync(std::unique_ptr const& decoder, DecodingOutput& output, DecodingInput const& input, - std::shared_ptr stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const maxBatchSize = input.batchSize; - - BufferManager::ITensorPtr finishedSum; - std::int32_t* finishedSumHost = nullptr; - if (input.sequenceLimitLength && output.finishReasons) - { - finishedSumHost = bufferCast(*output.finishedSum); - for (SizeType32 bi = 0; bi < maxBatchSize; ++bi) - { - finishedSumHost[bi] = 0; - } - } - - decoder->forwardAsync(output, input); - - if (finishedSumHost) - { - auto const numToFinish = output.finishReasons->getSize(); - TLLM_CUDA_CHECK(::cudaStreamSynchronize(stream->get())); - - SizeType32 finishedSum = 0; - for (SizeType32 bi = 0; bi < maxBatchSize; ++bi) - { - finishedSum += finishedSumHost[bi]; - } - return numToFinish == static_cast(finishedSum); - } - else - { - return false; - } -} - -void testDecoder(tensorrt_llm::DataType const dtype, SamplingConfig const& samplingConfig) -{ - SizeType32 constexpr tensorParallelism{1}; - SizeType32 constexpr pipelineParallelism{1}; - SizeType32 constexpr contextParallelism{1}; - SizeType32 constexpr localRank{0}; - WorldConfig const worldConfig{tensorParallelism, pipelineParallelism, contextParallelism, localRank}; - - SizeType32 constexpr vocabSize{51200}; - SizeType32 constexpr nbLayers{2}; - SizeType32 constexpr nbRnnLayers{0}; - SizeType32 constexpr nbHeads{16}; - SizeType32 constexpr hiddenSize{1024}; - SizeType32 constexpr batchSize{4}; - ModelConfig modelConfig{vocabSize, nbLayers + nbRnnLayers, nbLayers, nbRnnLayers, nbHeads, hiddenSize, dtype}; - modelConfig.useGptAttentionPlugin(false); - - SizeType32 constexpr maxInputLength{8}; - SizeType32 constexpr maxNewTokens{2}; - SizeType32 constexpr sinkTokenLength{0}; - auto constexpr maxSeqLength = maxInputLength + maxNewTokens; - - auto streamPtr = std::make_shared(); - BufferManager manager(streamPtr); - - // setup decoder - auto const beamWidth = samplingConfig.beamWidth; - - auto const decodingMode = beamWidth == 1 ? tle::DecodingMode::TopKTopP() : tle::DecodingMode::BeamSearch(); - - // create decoder - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - auto decoder = IGptDecoder::create( - decodingMode, modelConfig.getDataType(), batchSize, beamWidth, vocabSize, vocabSizePadded, streamPtr); - ASSERT_TRUE(static_cast(decoder)); - - auto batchSlots = getDefaultBatchSlots(batchSize); - decoder->setup(samplingConfig, batchSize, batchSlots); - - // set up inputs - std::vector> logitsVec; - for (auto i = 0; i < batchSize; ++i) - { - auto logits = manager.gpu(ITensor::makeShape({1, beamWidth, vocabSizePadded}), modelConfig.getDataType()); - manager.setZero(*logits); - logitsVec.push_back(std::move(logits)); - } - - int constexpr endId{50257}; - std::vector const endIdsVec(batchSize * beamWidth, endId); - auto endIds - = std::shared_ptr(manager.copyFrom(endIdsVec, ITensor::makeShape({batchSize, beamWidth}), MemoryType::kGPU)); - - DecodingInput inputs; - inputs.maxLength = maxInputLength; - inputs.maxAttentionWindow = maxSeqLength; - inputs.sinkTokenLength = sinkTokenLength; - inputs.batchSize = batchSize; - inputs.logitsVec = logitsVec; - inputs.endIds = endIds; - inputs.batchSlots = batchSlots; - - std::vector inputLengthsVec(batchSize * beamWidth, 0); - inputs.lengths = manager.copyFrom(inputLengthsVec, ITensor::makeShape({batchSize * beamWidth}), MemoryType::kGPU); - - std::vector sequenceLimitLengthsVec(batchSize, maxSeqLength); - inputs.sequenceLimitLength - = manager.copyFrom(sequenceLimitLengthsVec, ITensor::makeShape({batchSize}), MemoryType::kGPU); - - if (beamWidth > 1) - { - auto srcCacheIndirection = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); - manager.setZero(*srcCacheIndirection); - inputs.cacheIndirection = srcCacheIndirection; - } - - // set up outputs - auto outputIds = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); - manager.setZero(*outputIds); - auto gatheredOutputIds = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); - manager.setZero(*gatheredOutputIds); - DecodingOutput outputs{outputIds, gatheredOutputIds}; - auto newTokens - = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kINT32)); - manager.setZero(*newTokens); - outputs.newTokens = newTokens; - - std::vector sequenceLengthsVec(batchSize * beamWidth, maxInputLength); - outputs.lengths - = manager.copyFrom(sequenceLengthsVec, ITensor::makeShape({batchSize, beamWidth}), MemoryType::kGPU); - outputs.finishReasons = manager.gpu(ITensor::makeShape({batchSize, beamWidth}), - TRTDataType::value); - inputs.finishReasons = ITensor::view(outputs.finishReasons); - manager.setZero(*outputs.finishReasons); - outputs.finishedSum = BufferManager::pinnedPool(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); - auto finishedSumHost = bufferCast(*outputs.finishedSum); - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - finishedSumHost[bi] = -1; - } - - if (beamWidth > 1) - { - auto tgtCacheIndirection = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); - manager.setZero(*tgtCacheIndirection); - outputs.cacheIndirection = tgtCacheIndirection; - - auto cumLogProbs - = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kFLOAT)); - manager.setZero(*cumLogProbs); - outputs.cumLogProbs = cumLogProbs; - - auto parentIds = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); - manager.setZero(*parentIds); - outputs.parentIds = parentIds; - } - - // run decoder - EXPECT_FALSE(forwardAndSync(decoder, outputs, inputs, streamPtr)); - inputs.step += 1; - - { - SizeType32 finishedSum = 0; - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - finishedSum += finishedSumHost[bi]; - } - EXPECT_EQ(finishedSum, 0); - } - - // verify results - auto outputsIdsHost = manager.copyFrom(*outputs.ids, MemoryType::kCPU); - auto output = bufferCast(*outputsIdsHost); - manager.getStream().synchronize(); - - for (auto b = 0; b < batchSize; ++b) - { - for (auto bw = 0; bw < beamWidth; ++bw) - { - auto const result = (beamWidth == 1) ? 1023 : bw; - - bool anyMismatch = false; - for (auto i = 0; i < maxInputLength; ++i) - { - auto const outputIndex = tc::flat_index3(b, bw, i, beamWidth, maxSeqLength); - EXPECT_EQ(output[outputIndex], 0) << " b: " << b << " bw: " << bw << " i: " << i; - anyMismatch |= (output[outputIndex] != 0); - } - for (auto i = 0; i < maxNewTokens - 1; ++i) - { - auto const index = tc::flat_index3(b, bw, maxInputLength + i, beamWidth, maxSeqLength); - EXPECT_EQ(output[index], result) << " b: " << b << " bw: " << bw << " i: " << i; - anyMismatch |= (output[index] != result); - } - ASSERT_FALSE(anyMismatch); - } - } - - // run decoder again - EXPECT_TRUE(forwardAndSync(decoder, outputs, inputs, streamPtr)); - { - SizeType32 finishedSum = 0; - for (SizeType32 bi = 0; bi < batchSize; ++bi) - { - finishedSum += finishedSumHost[bi]; - } - EXPECT_EQ(finishedSum, outputs.finishReasons->getSize()); - } -} - -} // namespace - -class ParamTest : public ::testing::TestWithParam> -{ -}; - -TEST_P(ParamTest, Test) -{ - tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; - SizeType32 const beamWidth{std::get<1>(GetParam())}; - SamplingConfig const samplingConfig{beamWidth}; - - testDecoder(dtype, samplingConfig); -} - -INSTANTIATE_TEST_SUITE_P(DecoderTest, ParamTest, - testing::Combine( - testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), testing::Values(1, 3)), - [](testing::TestParamInfo const& info) - { - std::string name{std::get<0>(info.param) == tensorrt_llm::DataType::kFLOAT ? "Float" : "Half"}; - auto const beamWidth = std::get<1>(info.param); - name.append(beamWidth == 1 ? "Sampling" : "BeamWidth" + std::to_string(beamWidth)); - return name; - }); diff --git a/cpp/tests/unit_tests/runtime/samplingTest.cpp b/cpp/tests/unit_tests/runtime/samplingTest.cpp deleted file mode 100644 index bb93478cf8ff..000000000000 --- a/cpp/tests/unit_tests/runtime/samplingTest.cpp +++ /dev/null @@ -1,147 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/common/tllmDataType.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/layers/dynamicDecodeLayer.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/gptDecoder.h" - -#include - -using namespace tensorrt_llm::runtime; - -namespace tc = tensorrt_llm::common; -namespace tk = tensorrt_llm::kernels; -namespace tl = tensorrt_llm::layers; -namespace tle = tensorrt_llm::executor; - -class SamplingTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - void SetUp() override - { - mDeviceCount = tc::getDeviceCount(); - - if (mDeviceCount == 0) - GTEST_SKIP() << "No GPUs found"; - } - - void TearDown() override {} - - int mDeviceCount; -}; - -std::shared_ptr dynamicDecodeTest(std::shared_ptr manager, size_t vocabSize, - size_t vocabSizePadded, size_t batchSize, size_t beamWidth, int step, int ite, int maxInputLength, - size_t maxSeqLength, int localBatchSize, std::vector& cpuOutputIds, std::vector cpuLogits, - int noRepeatNgramSizeValue = 0) -{ - constexpr int endId = 1; - auto signedBatchSize = static_cast(batchSize); - auto signedBeamWidth = static_cast(beamWidth); - auto signedMaxSeqLength = static_cast(maxSeqLength); - cudaDeviceProp prop{}; - tc::check_cuda_error(cudaGetDeviceProperties(&prop, 0)); - - std::vector cpuEndIds(batchSize, endId); - std::vector cpuSequenceLengths(batchSize, maxInputLength); - std::vector cpuNoRepeatNgramSize(batchSize, noRepeatNgramSizeValue); - - tk::FinishedState::UnderlyingType* gpuFinished = nullptr; - - ITensor::SharedPtr gpuEndIds = manager->gpu(ITensor::makeShape({signedBatchSize}), tensorrt_llm::DataType::kINT32); - manager->copy(cpuEndIds.data(), *gpuEndIds, MemoryType::kCPU); - ITensor::SharedPtr gpuOutputIds = manager->gpu( - ITensor::makeShape({signedBatchSize, signedBeamWidth, signedMaxSeqLength}), tensorrt_llm::DataType::kINT32); - manager->copy(cpuOutputIds.data(), *gpuOutputIds, MemoryType::kCPU); - - auto const decodingMode = beamWidth == 1 ? tle::DecodingMode::TopKTopP() : tle::DecodingMode::BeamSearch(); - auto const decodingDomain = tensorrt_llm::layers::DecoderDomain(batchSize, beamWidth, vocabSize, vocabSizePadded); - auto ddLayer = tl::DynamicDecodeLayer(decodingMode, decodingDomain, manager); - - auto setupParams = std::make_shared(); - setupParams->banWordsParams = std::make_shared(); - setupParams->banWordsParams->noRepeatNgramSize = cpuNoRepeatNgramSize; - - setupParams->penaltyParams = std::make_shared(); - setupParams->decodingParams = std::make_shared(); - - auto batchSlots = getDefaultBatchSlots(batchSize); - auto workspace = std::make_shared( - manager, decodingDomain, TRTDataType::value, ddLayer.getWorkspaceSize()); - ddLayer.setup(batchSize, beamWidth, batchSlots, setupParams, workspace); - - auto forwardParams = std::make_shared(gpuEndIds, batchSlots, step, ite, localBatchSize); - auto logitsShape - = ITensor::makeShape({signedBatchSize, static_cast(beamWidth), static_cast(vocabSizePadded)}); - ITensor::SharedPtr inputLogits = manager->gpu(logitsShape, tensorrt_llm::DataType::kFLOAT); - forwardParams->logits = inputLogits; - manager->copy(cpuLogits.data(), *inputLogits, MemoryType::kCPU); - - forwardParams->banWordsInputs = std::make_shared(localBatchSize); - - forwardParams->stopCriteriaInputs = std::make_shared(localBatchSize); - - auto outputParams = std::make_shared(gpuOutputIds); - outputParams->sequenceLength = manager->gpu(ITensor::makeShape({signedBatchSize}), tensorrt_llm::DataType::kINT32); - manager->copy(cpuSequenceLengths.data(), *outputParams->sequenceLength.value(), MemoryType::kCPU); - outputParams->newTokens - = manager->gpu(ITensor::makeShape({signedBatchSize, signedBeamWidth}), tensorrt_llm::DataType::kINT32); - outputParams->finished = manager->gpu( - ITensor::makeShape({signedBatchSize, signedBeamWidth}), TRTDataType::value); - - ddLayer.forwardAsync(outputParams, forwardParams, workspace); - - return outputParams; -} - -TEST_F(SamplingTest, SamplingWithNoRepeatNGramSize) -{ - auto streamPtr = std::make_shared(); - auto manager = std::make_shared(streamPtr); - - constexpr size_t vocabSize{200}; - constexpr size_t vocabSizePadded{256}; - constexpr size_t batchSize{1}; - constexpr size_t beamWidth{1}; - constexpr int step{8}; - constexpr int ite{0}; - constexpr int maxInputLength{8}; - constexpr int maxSeqLength{9}; - constexpr int localBatchSize{batchSize}; - constexpr int noRepeatNgramSize{3}; - - std::vector cpuOutputIds(batchSize * beamWidth * maxSeqLength); - int ids[maxInputLength] = {10, 11, 12, 40, 41, 42, 40, 41}; - for (int i = 0; i < maxInputLength; i++) - { - cpuOutputIds[i] = ids[i]; - } - - std::vector cpuLogits(batchSize * beamWidth * vocabSizePadded, 1.0); - // We're setting 42 as favorite for greeding sampling but it should be banned because of no_repeat_ngram_size - // 43 is the expected fallback - cpuLogits[42] = 10.0; - cpuLogits[43] = 5.0; - - auto outputParams = dynamicDecodeTest(manager, vocabSize, vocabSizePadded, batchSize, beamWidth, step, ite, - maxInputLength, maxSeqLength, localBatchSize, cpuOutputIds, cpuLogits, noRepeatNgramSize); - - manager->copy(*outputParams->outputIds, cpuOutputIds.data(), MemoryType::kCPU); - EXPECT_EQ(cpuOutputIds[maxSeqLength - 1], 43); -} diff --git a/docs/source/blogs/tech_blog/blog11_GPT_OSS_Eagle3.md b/docs/source/blogs/tech_blog/blog11_GPT_OSS_Eagle3.md index 58b4230d157a..5e0b4c5f8042 100644 --- a/docs/source/blogs/tech_blog/blog11_GPT_OSS_Eagle3.md +++ b/docs/source/blogs/tech_blog/blog11_GPT_OSS_Eagle3.md @@ -90,7 +90,6 @@ speculative_config: speculative_model_dir: /config/models/eagle/ cuda_graph_config: max_batch_size: 10 -sampler_type: TorchSampler moe_config: backend: TRTLLM EOF @@ -99,7 +98,6 @@ EOF Notes: - Ensure your base model directory is `/config/models/gpt-oss-120b`. - Ensure your Eagle3 assets are present under `/config/models/eagle/`. -- On older releases (pre-1.1.0), replace `sampler_type: TorchSampler` with `use_torch_sampler: true`. ### Launch the Server (Eagle3 Speculative Decoding) diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 072cbd825c8a..8d443e07f558 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -217,7 +217,6 @@ unset or when the safety sanitizer rejects the runtime value. | `request_stats_max_iterations` | `Optional[int]` | `value` | | | | `return_perf_metrics` | `` | `value` | | | | `sampler_force_async_worker` | `` | `value` | | | -| `sampler_type` | `Union[str, tensorrt_llm.llmapi.llm_args.SamplerType]` | `categorical` | allowlist | `TRTLLMSampler`, `TorchSampler`, `auto` | | `scheduler_config.capacity_scheduler_policy` | `` | `categorical` | | `MAX_UTILIZATION`, `GUARANTEED_NO_EVICT`, `STATIC_BATCH` | | `scheduler_config.context_chunking_policy` | `Optional[tensorrt_llm.llmapi.llm_args.ContextChunkingPolicy]` | `categorical` | | `FIRST_COME_FIRST_SERVED`, `EQUAL_PROGRESS`, `FORCE_CHUNK` | | `scheduler_config.dynamic_batch_config.dynamic_batch_moving_average_window` | `` | `value` | | | diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index d9d3fb6f8219..1494eba98a95 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -17,27 +17,8 @@ The PyTorch backend supports a wide variety of features, listed below: ## General usage -There are two sampling backends available. - -* Torch Sampler -* TRTLLM Sampler (deprecated) - -Torch Sampler is used by default and supports a superset of features of TRTLLM Sampler. TRTLLM Sampler will be removed in release 1.4. -One can specify which sampler to use explicitly with: - -```python -from tensorrt_llm import LLM - -# Chooses TorchSampler explicitly -llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8', - sampler_type="TorchSampler") - -# Chooses TRTLLMSampler explicitly -llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8', - sampler_type="TRTLLMSampler") -``` - -By default, the sampling backend is chosen to be `auto`. This will use Torch Sampler for all requests. +Torch Sampler is the only sampling backend and is used for all requests; there +is nothing to configure. Here is an example to run a model with basic usage of sampling parameters. This example prepares two identical prompts which will give different results due to the sampling parameters chosen: diff --git a/examples/auto_deploy/paragraf/create_standalone_package.py b/examples/auto_deploy/paragraf/create_standalone_package.py index d2e75e5cbf39..a3809e46ad0b 100644 --- a/examples/auto_deploy/paragraf/create_standalone_package.py +++ b/examples/auto_deploy/paragraf/create_standalone_package.py @@ -173,7 +173,6 @@ "test_ad_build_small_single.py": "test_paragraf_trtllm_build_small_single.py", "test_ad_guided_decoding_regex.py": "test_paragraf_trtllm_guided_decoding_regex.py", "test_ad_trtllm_bench.py": "test_paragraf_trtllm_bench.py", - "test_ad_trtllm_sampler.py": "test_paragraf_trtllm_sampler.py", "test_ad_trtllm_serve.py": "test_paragraf_trtllm_serve.py", } SOURCE_TEST_NAMES_BY_GENERATED_NAME = { diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 9bdf74f6b1c6..b8efc9a20281 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -100,9 +100,6 @@ def add_llm_args(parser): parser.add_argument('--attention_dp_batching_wait_iters', type=int, default=0) - parser.add_argument('--sampler_type', - default="auto", - choices=["auto", "TorchSampler", "TRTLLMSampler"]) parser.add_argument('--tp_size', type=int, default=1) parser.add_argument('--pp_size', type=int, default=1) parser.add_argument('--orchestrator_type', @@ -377,7 +374,6 @@ def setup_llm(args, **kwargs): args.use_piecewise_cuda_graph) if args.use_torch_compile else None, moe_config=MoeConfig(backend=args.moe_backend, use_low_precision_moe_combine=args.use_low_precision_moe_combine, load_balancer=args.moe_load_balancer_config), - sampler_type=args.sampler_type, max_seq_len=args.max_seq_len, max_batch_size=args.max_batch_size, max_num_tokens=args.max_num_tokens, diff --git a/legacy-files.txt b/legacy-files.txt index 73fe3fee5899..0f4694eb5d55 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -298,7 +298,6 @@ tensorrt_llm/_torch/pyexecutor/handle_logits.py tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py tensorrt_llm/_torch/pyexecutor/llm_request.py -tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py tensorrt_llm/_torch/pyexecutor/model_engine.py tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -611,7 +610,6 @@ tests/unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py tests/unittest/_torch/sampler/test_beam_search.py tests/unittest/_torch/sampler/test_best_of_n.py -tests/unittest/_torch/sampler/test_trtllm_sampler.py tests/unittest/_torch/speculative/test_eagle3.py tests/unittest/_torch/test_connector.py tests/unittest/_torch/test_torch_multi_arange.py diff --git a/pyproject.toml b/pyproject.toml index f8e1447ca03b..d57d9c1d2d63 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -355,7 +355,6 @@ exclude = [ "tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py", "tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py", "tensorrt_llm/_torch/pyexecutor/llm_request.py", - "tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py", "tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py", "tensorrt_llm/_torch/pyexecutor/model_engine.py", "tensorrt_llm/_torch/pyexecutor/model_loader.py", @@ -668,7 +667,6 @@ exclude = [ "tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py", "tests/unittest/_torch/sampler/test_beam_search.py", "tests/unittest/_torch/sampler/test_best_of_n.py", - "tests/unittest/_torch/sampler/test_trtllm_sampler.py", "tests/unittest/_torch/speculative/test_eagle3.py", "tests/unittest/_torch/test_connector.py", "tests/unittest/_torch/test_torch_multi_arange.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 612b315564ca..a7d38f4c7beb 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -315,7 +315,6 @@ include = [ "tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py", "tensorrt_llm/_torch/pyexecutor/layerwise_nvtx_marker.py", "tensorrt_llm/_torch/pyexecutor/llm_request.py", - "tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py", "tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py", "tensorrt_llm/_torch/pyexecutor/model_engine.py", "tensorrt_llm/_torch/pyexecutor/model_loader.py", @@ -628,7 +627,6 @@ include = [ "tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py", "tests/unittest/_torch/sampler/test_beam_search.py", "tests/unittest/_torch/sampler/test_best_of_n.py", - "tests/unittest/_torch/sampler/test_trtllm_sampler.py", "tests/unittest/_torch/speculative/test_eagle3.py", "tests/unittest/_torch/test_connector.py", "tests/unittest/_torch/test_torch_multi_arange.py", diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 1ed345a1914f..bbd4009b6d7e 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -22,7 +22,6 @@ from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.distributed import Distributed -from tensorrt_llm._torch.pyexecutor._util import get_decoding_mode from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.guided_decoder import GuidedDecoder from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import ( @@ -39,7 +38,7 @@ ResourceManager, ResourceManagerType, ) -from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler, TRTLLMSampler +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler from tensorrt_llm._torch.pyexecutor.scheduler import ( BindCapacityScheduler, BindMicroBatchScheduler, @@ -51,7 +50,7 @@ from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler from tensorrt_llm._utils import get_free_port, mpi_rank, mpi_world_size, nvtx_range from tensorrt_llm.inputs.multimodal import MultimodalRuntimeData, check_mm_embed_cumsum_if_needed -from tensorrt_llm.llmapi.llm_args import ContextChunkingPolicy, MultimodalConfig, SamplerType +from tensorrt_llm.llmapi.llm_args import ContextChunkingPolicy, MultimodalConfig from tensorrt_llm.llmapi.tokenizer import TokenizerBase from tensorrt_llm.mapping import Mapping @@ -1079,17 +1078,6 @@ def forward( return outputs -class TRTLLMSamplerModelConfig: - def __init__(self, vocab_size_padded: int): - self.config = SimpleNamespace() - self.config.vocab_size = vocab_size_padded - - # Initialized to dummy values as they are not used in the C++ code underlying TRTLLMSampler. - self.config.num_hidden_layers = 42 - self.config.hidden_size = 42 - self.config.num_attention_heads = 42 - - def instantiate_sampler( ad_config: LlmArgs, max_num_sequences: int, @@ -1115,42 +1103,16 @@ def instantiate_sampler( ) return SpecSampler(sampler_args) - sampler_type = ad_config.sampler_type - if sampler_type == SamplerType.auto: - sampler_type = SamplerType.TorchSampler - - if sampler_type == SamplerType.TorchSampler: - # Regular TorchSampler for non-speculative decoding. - sampler_args = TorchSampler.Args( - max_seq_len=ad_config.max_seq_len, - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - max_num_sequences=max_num_sequences, - max_beam_width=ad_config.max_beam_width, - disable_overlap_scheduler=ad_config.disable_overlap_scheduler, - ) - sampler = TorchSampler(sampler_args) - - elif sampler_type == SamplerType.TRTLLMSampler: - vocab_size_padded: int = engine.cache_seq_interface.info.vocab_size_padded - sampler_model_config = TRTLLMSamplerModelConfig(vocab_size_padded) - decoding_mode = get_decoding_mode(ad_config.decoding_config, ad_config.max_beam_width) - sampler = TRTLLMSampler( - model=sampler_model_config, - model_dtype=torch.bfloat16, # hardcoded as bfloat16; does not seem necessary in C++ code. - mapping=dist_mapping, - decoding_mode=decoding_mode, - disable_overlap_scheduler=ad_config.disable_overlap_scheduler, - max_seq_len=ad_config.max_seq_len, - max_batch_size=ad_config.max_batch_size, - max_beam_width=ad_config.max_beam_width, - decoding_config=ad_config.decoding_config, - kv_cache_config=ad_config.kv_cache_config, - ) - else: - raise ValueError(f"Sampler type {sampler_type} is not supported.") - - return sampler + # Regular TorchSampler for non-speculative decoding. + sampler_args = TorchSampler.Args( + max_seq_len=ad_config.max_seq_len, + max_draft_len=max_draft_len, + max_total_draft_tokens=max_total_draft_tokens, + max_num_sequences=max_num_sequences, + max_beam_width=ad_config.max_beam_width, + disable_overlap_scheduler=ad_config.disable_overlap_scheduler, + ) + return TorchSampler(sampler_args) def create_autodeploy_executor( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 5835d677b6f0..c0e2d07f9009 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -15,7 +15,7 @@ import copy import dataclasses import os -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import torch @@ -24,16 +24,14 @@ from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, is_sm_100f, prefer_pinned, str_dtype_to_binding, torch_dtype_to_str) -from tensorrt_llm.bindings.executor import DecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams # isort: off from tensorrt_llm.llmapi.llm_args import ( CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, - MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SamplerType, - SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, - WaitingQueuePolicy) + MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SchedulerConfig, + SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy) # isort: on from tensorrt_llm._torch.peft.lora.config import ( LoraConfig, get_default_trtllm_modules_to_hf_modules) @@ -72,8 +70,7 @@ from .resource_manager import (KVCacheCompressionManager, KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) -from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, - TRTLLMSampler) +from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, KVCacheV2Scheduler, MultimodalEagerEncoderScheduler, MultimodalScheduler, SimpleScheduler, @@ -3335,11 +3332,8 @@ def instantiate_sampler( *, max_batch_size: int, max_beam_width: int, - max_seq_len: int, mm_encoder_only: bool, speculative_config: SpeculativeConfig, - decoding_config: trtllm.DecodingConfig, - kv_cache_config: KvCacheConfig, max_num_sequences: Optional[int] = None, ): enable_async_worker = (confidential_compute_enabled() @@ -3357,8 +3351,6 @@ def instantiate_sampler( enable_speculative_beam_history_d2h, max_num_sequences=max_num_sequences, ) - decoding_mode = get_decoding_mode(decoding_config=decoding_config, - max_beam_width=max_beam_width) if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): return get_spec_decoder(sampler_args, engine.spec_config) @@ -3366,52 +3358,12 @@ def instantiate_sampler( if mm_encoder_only: # NOTE: handle model outputs specially for mm encoder executor/engine return EarlyStopWithMMResult() - if llm_args.sampler_type == SamplerType.TRTLLMSampler: - logger.warning( - "TRTLLMSampler is deprecated and will be removed in release 1.4. Please use TorchSampler instead." - ) - logger.debug(f"DecodingMode: {decoding_mode.name}") - return TRTLLMSampler(engine.model, - engine.dtype, - mapping, - decoding_mode, - llm_args.disable_overlap_scheduler, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - decoding_config=decoding_config, - kv_cache_config=kv_cache_config, - enable_async_worker=enable_async_worker, - max_num_sequences=max_num_sequences) if not engine.model.model_config.is_generation: # NOTE: choose sampler based on model type return EarlyStopSampler() return TorchSampler(sampler_args) -def get_decoding_mode( - decoding_config: trtllm.DecodingConfig, - max_beam_width: int, -) -> DecodingMode: - '''This implementation is based off trtGptModelInflightBatching.cpp getDecodingMode().''' - if decoding_config and decoding_config.decoding_mode and not decoding_config.decoding_mode.isAuto( - ): - decoding_mode = decoding_config.decoding_mode - elif max_beam_width == 1: - decoding_mode = DecodingMode.TopKTopP() - else: - decoding_mode = DecodingMode.BeamSearch() - - # Override decoding mode when beam width is one - if max_beam_width == 1 and decoding_mode.isBeamSearch(): - logger.warning( - "Beam width is set to 1, but decoding mode is BeamSearch. Overwriting decoding mode to TopKTopP." - ) - decoding_mode = DecodingMode.TopKTopP() - - return decoding_mode - - _ATTN_MODULES = frozenset({ "attn_q", "attn_k", @@ -3582,7 +3534,7 @@ def _adjust_torch_mem_fraction(): torch.cuda.set_per_process_memory_fraction(mem_torch_fraction) -def validate_feature_combination(llm_args, model_engine, sampler_type): +def validate_feature_combination(llm_args, model_engine): # Validate the flags for features' combination compression_config = llm_args.kv_cache_compression_config if compression_config is not None: @@ -3605,8 +3557,6 @@ def init_feature_status(llm_args) -> Dict[str, bool]: "mtp", "eagle3_one_model", "eagle3_two_model", - "torch_sampler", - "trtllm_sampler", "kv_cache_reuse", "slide_window_attention", "guided_decoding", @@ -3627,11 +3577,6 @@ def init_feature_status(llm_args) -> Dict[str, bool]: feature_status["eagle3_two_model"] = ( isinstance(llm_args.speculative_config, EagleDecodingConfig) and not llm_args.speculative_config.eagle3_one_model) - feature_status[ - "torch_sampler"] = sampler_type == SamplerType.TorchSampler - feature_status[ - "trtllm_sampler"] = sampler_type == SamplerType.TRTLLMSampler - feature_status[ "kv_cache_reuse"] = llm_args.kv_cache_config is not None and llm_args.kv_cache_config.enable_block_reuse feature_status["slide_window_attention"] = ( @@ -3649,29 +3594,9 @@ def init_feature_status(llm_args) -> Dict[str, bool]: feature_status: Dict[str, bool] = init_feature_status(llm_args) - ERR_MSG_TMPL = "{feature1} and {feature2} enabled together is not supported yet." - - CONFLICT_RULES = [ - { - "features": ["trtllm_sampler", "mtp"], - "message": - ERR_MSG_TMPL.format(feature1="trtllm_sampler", feature2="mtp") + - " Please use sampler type auto instead." - }, - { - "features": ["trtllm_sampler", "eagle3_one_model"], - "message": - ERR_MSG_TMPL.format(feature1="trtllm_sampler", - feature2="eagle3_one_model") + - " Please use sampler type auto instead." - }, - { - "features": ["trtllm_sampler", "eagle3_two_model"], - "message": - ERR_MSG_TMPL.format(feature1="trtllm_sampler", - feature2="eagle3_two_model") + - " Please use sampler type auto instead." - }, + # Kept as an extension point; there are currently no conflicting feature + # combinations to reject. + CONFLICT_RULES: list[dict[str, Any]] = [ # Add new conflict rules here in the future ] for rule in CONFLICT_RULES: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 696eceed0e57..4d2106c65c6e 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1038,8 +1038,7 @@ def __init__( self._initialize_execution_state(seq_slot=seq_slot, orig_prompt_len=self.orig_prompt_len) - # TODO: remove this when use DynamicDecodeOp in pytorch flow. - # currently, keep py_stop_words_list as python list, rather than tensor. + # Keep py_stop_words_list as a python list, rather than a tensor. self.py_stop_words_list = stop_words_list self.py_logprobs_mode = LogprobMode( diff --git a/tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py b/tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py deleted file mode 100644 index 5cb427f00b29..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/make_decoding_batch_input_output.py +++ /dev/null @@ -1,72 +0,0 @@ -from dataclasses import dataclass -from typing import List - -import torch - -from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.bindings.internal.batch_manager import DecoderInputBuffers -from tensorrt_llm.bindings.internal.runtime import DecoderState - - -@dataclass -class MakeDecodingBatchInputOutput: - """Python implementation of MakeDecodingBatchInputOutput algorithm. - - This class is responsible for creating decoder batch inputs and outputs for the decoding process. - It handles both context and generation requests, managing their logits and batch slots. - """ - - @torch.inference_mode() - @nvtx_range("make_decoding_batch_input_output") - def __call__( - self, - decoder_input_buffers: DecoderInputBuffers, - decoder_state: DecoderState, - scheduled_requests, - logits: torch.Tensor, - beam_width: int, - num_context_logits_prefix_sum: List[int], - ): - """Create decoder batch inputs and outputs for the given requests. - - Args: - decoder_input_buffers: Decoder input buffers - decoder_state: Current decoder state - scheduled_requests: Scheduled requests - logits: Logits tensor - beam_width: Beam width - num_context_logits_prefix_sum: Number of context logits prefix sum - """ - # In order to make a decoding_input assuming no drafting, we need: - # 1. logits_vec = [[logits_slice of each active slot]] - # 2. batch_slots = [[active_slots]] - # 3. generation_steps = [decoding_iters] - - active_slots = [[]] - generation_steps = [] - logits_vec = [[]] - for i, r in enumerate( - scheduled_requests.context_requests_last_chunk, - start=len(scheduled_requests.context_requests_chunking)): - active_slots[0].append(r.py_seq_slot) - generation_steps.append(r.decoding_iter) - logits_vec[0].append( - logits[num_context_logits_prefix_sum[i]: - num_context_logits_prefix_sum[i + 1]].unsqueeze(0)) - - logits_index = num_context_logits_prefix_sum[-1] - for i, r in enumerate(scheduled_requests.generation_requests): - if r.is_generation_in_progress_state: - active_slots[0].append(r.py_seq_slot) - generation_steps.append(r.decoding_iter) - logits_vec[0].append( - logits.narrow(dim=0, - start=logits_index + i * beam_width, - length=beam_width).unsqueeze(0)) - - decoder_state.generation_steps = generation_steps - decoder_input_buffers.forward_batch_slots = [ - torch.tensor(active_slots[0], dtype=torch.int32) - ] - decoder_input_buffers.logits = logits_vec - decoder_input_buffers.max_decoder_steps = 1 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c3cd9e97ea37..cde82f14b45e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -95,7 +95,7 @@ from .resource_manager import (NoFreeSlotsError, ResourceManager, ResourceManagerType, request_context) from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, - SampleStateTensors, TRTLLMSampler) + SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, WaitingQueue, create_waiting_queue) @@ -692,15 +692,6 @@ def __init__( "TRTLLM_PP_MULTI_STREAM_SAMPLE", "1") == "1" self.sample_stream = torch.cuda.Stream() self.finish_sample_event = torch.cuda.Event() - if (self.dist.pp_size > 1 and self.pp_multi_stream_sample - and isinstance(self.sampler, TRTLLMSampler)): - # TRTLLM sampler uses default stream for store and algorithms. - # To enable multi-stream sampling, we need to re-initialize - # the sampler store and algorithms on the sample stream. - with torch.cuda.stream(self.sample_stream): - self.sampler._initialize_store() - self.sampler._instantiate_algorithms() - # Set of request IDs that are currently in flight across all micro batches # or waiting for synchronized PP resource teardown. The scheduler avoids # these requests until their prior execution state is safe to reuse. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index c4bc59ec28fe..039f81b06ff1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -206,19 +206,6 @@ def _get_mapping(_mapping: Mapping) -> Mapping: return mapping -def update_sampler_max_seq_len(max_seq_len, sampler): - # Originally, TRTLLMSampler is constructed with executor_config, but - # _create_kv_cache_manager (via build_managers) may later overwrite executor_config.max_seq_len. - # Because TRTLLMSampler.sample_async still needs the updated limit and executor_config is - # deprecated inside TRTLLMSampler, keep TRTLLMSampler.max_seq_len updated with - # with executor_config.max_seq_len. - from .sampler import TRTLLMSampler - - if isinstance(sampler, TRTLLMSampler): - assert hasattr(sampler, "max_seq_len") - sampler.max_seq_len = max_seq_len - - def _extend_full_attention_windows_for_spec_decode( kv_cache_config: KvCacheConfig, spec_config: Optional[SpeculativeConfig], @@ -382,8 +369,6 @@ def create_py_executor( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False - decoding_config = llm_args.decoding_config - # The tokenizer is stripped from MPI kwargs in proxy.py to avoid pickle # failures with trust_remote_code models. Reload it from the checkpoint # when guided decoding needs it. @@ -590,7 +575,7 @@ def allocation_scope(current_stage: ExecutorMemoryType): model_weights_restore_mode=model_weights_restore_mode, ) - validate_feature_combination(llm_args, model_engine, llm_args.sampler_type) + validate_feature_combination(llm_args, model_engine) calibrator = get_calibrator() layer_wise_benchmarks_config = llm_args.layer_wise_benchmarks_config @@ -800,11 +785,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): mapping, max_batch_size=max_batch_size, max_beam_width=max_beam_width, - max_seq_len=max_seq_len, mm_encoder_only=mm_encoder_only, speculative_config=spec_config, - decoding_config=decoding_config, - kv_cache_config=kv_cache_config, max_num_sequences=max_num_seq_slots, ) logger.info(f"Using Sampler: {type(sampler).__name__}") @@ -933,10 +915,6 @@ def allocation_scope(current_stage: ExecutorMemoryType): ExecutorMemoryType.INIT_KV_CACHE if estimating_kv_cache else ExecutorMemoryType.KV_CACHE): kv_cache_creator.build_managers(resources, estimating_kv_cache) - # Originally, max_seq_len might be mutated inside build_managers as field of executor config. - # Since now, we are changing kv_cache_creator._max_seq_len instead. Restore max_seq_len here. - max_seq_len = kv_cache_creator._max_seq_len - update_sampler_max_seq_len(max_seq_len, sampler) # DWDP setup: MNNVL handle exchange + composite VA weight buffer + # weight manager + MoE backend fixup (single entry point). @@ -1032,10 +1010,6 @@ def allocation_scope(current_stage: ExecutorMemoryType): # the original value before creating the final KV cache. kv_cache_creator._max_seq_len = model_engine_max_seq_len kv_cache_creator.build_managers(resources, False) - # Originally, max_seq_len might be mutated inside build_managers as field of executor config. - # Since now, we are changing kv_cache_creator._max_seq_len instead. Restore max_seq_len here. - max_seq_len = kv_cache_creator._max_seq_len - update_sampler_max_seq_len(max_seq_len, sampler) with allocation_scope(ExecutorMemoryType.EXTRA_RESOURCES): diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/__init__.py b/tensorrt_llm/_torch/pyexecutor/sampler/__init__.py index 39989d8401de..976e2f79b239 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/__init__.py @@ -14,7 +14,7 @@ """Sampler package. -The upper-level orchestration (``Sampler`` / ``TorchSampler`` / ``TRTLLMSampler``) +The upper-level orchestration (``Sampler`` / ``TorchSampler``) lives in ``sampler.py`` and depends on operation-level APIs in ``sampler_strategy.py``. Implementation-specific kernel providers (FlashInfer, vanilla/PyTorch, TRT-LLM ops) live under ``ops/`` and are selected diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 680e31d256ad..f90046abade3 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -17,20 +17,12 @@ Holds the :class:`Sampler` ABC and its ``SampleState`` types, the trivial early-stop samplers (:class:`EarlyStopSampler` for non-generation models, :class:`EarlyStopWithMMResult` for the multimodal-encoder-only engine), and -the two real implementations: - -* :class:`TorchSampler` -- the PyTorch path. It owns no sampling logic of its - own beyond batching and orchestration: each feature (beam search, penalties, - token bans, top-p decay, finish reasons, log-probs, two-model speculation, - seeds) lives in its own module and is held here as a handler, driven through - ``setup_sampler_step`` / ``sample_async`` / ``update_requests``. -* :class:`TRTLLMSampler` -- the C++ decoder path. - -NOTE: ``TRTLLMSampler`` is deprecated and slated for removal (``_util.py`` -warns that it goes away in release 1.4). Dropping it -- together with -:class:`Algorithms` and its ``SampleState*`` types -- leaves only the ABC, the -state types and ``TorchSampler``, which do not need to share a file; this -module is worth reorganizing at that point. +:class:`TorchSampler` -- the PyTorch sampling path. ``TorchSampler`` owns no +sampling logic of its own beyond batching and orchestration: each feature +(beam search, penalties, token bans, top-p decay, finish reasons, log-probs, +two-model speculation, seeds) lives in its own module and is held here as a +handler, driven through ``setup_sampler_step`` / ``sample_async`` / +``update_requests``. """ import sys @@ -55,39 +47,14 @@ import torch from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE -from tensorrt_llm._torch.pyexecutor.make_decoding_batch_input_output import ( - MakeDecodingBatchInputOutput, -) -from tensorrt_llm._utils import mpi_disabled, nvtx_range, prefer_pinned, torch_dtype_to_binding -from tensorrt_llm.bindings import ( - CudaStream, - DataType, - ModelConfig, - SamplingConfigVector, - WorldConfig, - make_sampling_config, -) -from tensorrt_llm.bindings.executor import DecodingConfig, DecodingMode, FinishReason -from tensorrt_llm.bindings.internal.algorithms import CreateNewDecoderRequests -from tensorrt_llm.bindings.internal.batch_manager import ( - DecoderInputBuffers, - add_new_tokens_to_requests, - make_decoding_batch_input, -) -from tensorrt_llm.bindings.internal.runtime import ( - BufferManager, - CudaEvent, - DecoderState, - GptDecoderBatched, -) +from tensorrt_llm._utils import nvtx_range, prefer_pinned +from tensorrt_llm.bindings.executor import FinishReason +from tensorrt_llm.bindings.internal.batch_manager import add_new_tokens_to_requests from tensorrt_llm.executor.result import Logprob -from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping from tensorrt_llm.sampling_params import SamplingParams from ...utils import torch_multi_arange -from ..finish_reason import FinishedState from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests @@ -148,7 +115,7 @@ if TYPE_CHECKING: from transformers import PretrainedConfig - from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM + from tensorrt_llm._torch.models.modeling_utils import DecoderModel # Type-only: importing the speculative package at module level would # re-create the import cycle sampler.sampler -> speculative -> @@ -482,7 +449,7 @@ class Store: new_tokens: torch.Tensor """Device tensor containing latest sampled tokens. - Shape: See cpp DecoderState.getAllNewTokens(). + Shape: ``NEW_TOKENS_SHAPE`` -- (max_tokens, max_num_sequences, max_beam_width). """ beam_search_store: "BeamSearchStore | None" = None """Holds data related to beam search.""" @@ -1945,8 +1912,7 @@ def _select_generated_logits( # another request's rows. logits.view() succeeds for any shape whose # element count divides, so that is silent corruption rather than an # error. Match the layout here and slice down to the per-iteration - # width where the beams are actually consumed. TRTLLMSampler already - # offsets by the static width for the same reason. + # width where the beams are actually consumed. # NB: context requests do not have multiple beams yet, hence the 1s. req_num_beams_list = [1] * len(finished_context_requests) + [ req.py_beam_width for req in scheduled_requests.generation_requests @@ -2323,569 +2289,3 @@ def should_provide_draft_probs(self, request: LlmRequest) -> bool: use_beam_search=self._use_beam_search, min_p=min_p, ) - - -class Algorithms: - def defined_algorithms(self) -> list[str]: - return [attr for attr in dir(self) if not attr.startswith("__")] - - def __repr__(self) -> str: - algs = self.defined_algorithms() - return f"Algs({', '.join(algs)})" - - -@dataclass(kw_only=True) -class SampleStateTensorsHostTRTLLM(SampleStateTensors): - finished_sum: torch.Tensor - finish_reasons: torch.Tensor - sequence_lengths: torch.Tensor - cum_log_probs: torch.Tensor | None = None - gathered_ids: torch.Tensor | None = None - - -@dataclass(kw_only=True) -class SampleStateTRTLLM(SampleState[SampleStateTensorsHostTRTLLM, SampleStateTensors]): - finalize_events: dict[int, CudaEvent] | None = None - """`Optional` to accommodate `_forward_step_inter_pp` which creates a `SampleState` without `finalize_events`""" - - -class TRTLLMSampler(Sampler[SampleStateTRTLLM], AsyncWorkerMixin): - MAX_DECODING_TOKENS = 1 # It must be 1 when not in speculative decoding - SampleState = SampleStateTRTLLM - - @override - def is_generation_model(self) -> bool: - return True - - def __init__( - self, - model: "DecoderModelForCausalLM[_ModelType, _ConfigType]", - model_dtype: torch.dtype, - mapping: Mapping, - decoding_mode: DecodingMode, - disable_overlap_scheduler: bool, - max_seq_len: int, - max_batch_size: int, - max_beam_width: int, - decoding_config: Optional[DecodingConfig] = None, - kv_cache_config: Optional[KvCacheConfig] = None, - enable_async_worker: bool = False, - max_num_sequences: Optional[int] = None, - ): - assert model.config is not None - vocab_size = model.config.vocab_size - num_hidden_layers = model.config.num_hidden_layers - hidden_size = model.config.hidden_size - num_heads = model.config.num_attention_heads - - self.model_datatype = torch_dtype_to_binding(model_dtype) - self.logits_datatype = DataType.FLOAT - self.decoding_mode = decoding_mode - self.decoding_config = decoding_config if decoding_config else DecodingConfig(decoding_mode) - max_attn_window = kv_cache_config.max_attention_window # type: ignore - self.max_seq_len = max_seq_len - self.max_attention_window = ( - max(max_attn_window) if max_attn_window is not None else max_seq_len - ) - self.max_batch_size = max_batch_size - self.max_beam_width = max_beam_width - self.max_seq_idle_microseconds = 180 * 1000 * 1000 - self.is_trt_overlap = not disable_overlap_scheduler - self.num_micro_batches = ( - mapping.pp_size if mapping.pp_size > 1 else (2 if self.is_trt_overlap else 1) - ) - # Decoder state is indexed by sequence slot and must match the - # executor's SeqSlotManager. The fallback preserves the established - # sizing for direct callers outside the PyExecutor creator. - self.max_num_sequences = ( - max_num_sequences if max_num_sequences is not None else mapping.pp_size * max_batch_size - ) - self.micro_batch_idx = 0 - - if mpi_disabled(): - self.world_config = WorldConfig( - mapping.tp_size, - mapping.pp_size, - mapping.cp_size, - rank=mapping.rank, - gpus_per_node=mapping.gpus_per_node, - ) - else: - self.world_config = WorldConfig.mpi( - mapping.gpus_per_node, mapping.tp_size, mapping.pp_size - ) - self.model_config = ModelConfig( - vocab_size, - num_hidden_layers, - num_hidden_layers, - 0, - num_heads, - hidden_size, - self.model_datatype, - ) - - self._initialize_store() - self._instantiate_algorithms() - - self._async_worker_init(enable_async_worker) - - def _initialize_store(self) -> None: - torch_stream = torch.cuda.current_stream().cuda_stream - cuda_stream = CudaStream(torch_stream) - buffer_manager = BufferManager(stream=torch_stream) - - self.store = { - "torch_stream": torch_stream, - "cuda_stream": cuda_stream, - "buffer_manager": buffer_manager, - "decoder_input_buffers": [ - DecoderInputBuffers(self.max_batch_size, self.MAX_DECODING_TOKENS, buffer_manager) - for _ in range(self.num_micro_batches) - ], - "sequence_lengths_host": torch.empty( - ( - self.max_num_sequences, - self.max_beam_width, - ), - dtype=torch.int, - ), - "decoder_state": DecoderState(), - } - - cast(DecoderState, self.store["decoder_state"]).setup( - max_num_sequences=self.max_num_sequences, - max_beam_width=self.max_beam_width, - max_attention_window=self.max_attention_window, - sink_token_length=0, - max_sequence_length=self.max_seq_len, - dtype=self.logits_datatype, - model_config=self.model_config, - world_config=self.world_config, - buffer_manager=buffer_manager, - ) - - def _instantiate_algorithms(self) -> None: - self.algs = Algorithms() - self.algs.decoder = GptDecoderBatched(stream=self.store["torch_stream"]) # type: ignore - self.algs.decoder.setup( # type: ignore - mode=self.decoding_mode, - max_num_sequences=self.max_num_sequences, - max_beam_width=self.max_beam_width, - dtype=self.logits_datatype, - model_config=self.model_config, - world_config=self.world_config, - ) - self.algs.create_new_decoder_requests = CreateNewDecoderRequests( # type: ignore - speculative_decoding_fast_logits=False, - is_leader_in_orch_mode=False, - is_normalize_log_probs=False, - ) - self.algs.make_decoding_batch_input_output = MakeDecodingBatchInputOutput() # type: ignore - - @torch.inference_mode() - @nvtx_range("setup_sampler_step") - def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: - batch_slots, sampling_configs, lookahead_prompt, lookahead_algo_configs = ( - self.algs.create_new_decoder_requests( # type: ignore - self.model_config, - self.world_config, - self.decoding_config, - scheduled_requests.context_requests, - self.logits_datatype, - self.store["decoder_input_buffers"][self.micro_batch_idx], # type: ignore - self.store["decoder_state"], - self.store["cuda_stream"], - self.algs.decoder.decoder_stream, # type: ignore - self.max_seq_len, - self.beam_width(scheduled_requests.context_requests), - ) - ) - - local_batch_size = len(batch_slots) - if local_batch_size > 0: - sampling_config = make_sampling_config(sampling_configs) - self.algs.decoder.underlying_decoder().setup( # type: ignore - sampling_config, - local_batch_size, - batch_slots, - self.store["decoder_state"].joint_decoding_output, # type: ignore - self.model_config.data_type, - lookahead_prompt, - lookahead_algo_configs, - ) - - adp = [r for r in scheduled_requests.generation_requests if r.is_attention_dp_dummy] - batch_size = len(adp) - if batch_size == 0: - return - config = make_sampling_config(cast(SamplingConfigVector, [r.sampling_config for r in adp])) - slots = torch.tensor([r.py_seq_slot for r in adp], dtype=torch.int32) - self.algs.decoder.underlying_decoder().setup(config, batch_size, slots) # type: ignore - - def get_cache_indirection(self) -> torch.Tensor | None: - return self.store["decoder_state"].cache_indirection_output # type: ignore - - def _update_cache_indirection_buffer(self, scheduled_requests: ScheduledRequests) -> None: - # Copy cache indirection output to input - for request in scheduled_requests.generation_requests: - self.store["decoder_state"].cache_indirection_input[request.py_seq_slot].copy_( # type: ignore - self.store["decoder_state"].cache_indirection_output[request.py_seq_slot], # type: ignore - non_blocking=True, - ) - - @override - def validate_request(self, request: LlmRequest) -> None: - if ( - self.max_batch_size > 1 - and self.beam_width([request]) > 1 - and request.py_return_log_probs - ): - raise ValueError("Beam search only supports logprobs when batch size is 1") - - @torch.inference_mode() - @nvtx_range("sample_async") - @override - def sample_async( - self, - scheduled_requests: ScheduledRequests, - model_outputs: dict[str, Any], - num_context_logits_prefix_sum: list[int], - resource_manager: Optional[ResourceManager] = None, - ) -> SampleStateTRTLLM: - batch_size = scheduled_requests.batch_size - beam_width = self.beam_width(scheduled_requests.all_requests()) - assert not ( - batch_size > 1 - and beam_width > 1 - and any(request.py_return_log_probs for request in scheduled_requests.all_requests()) - ), "Beam search only supports logprobs when batch size is 1" - - self.setup_sampler_step(scheduled_requests) - - # For beam search, cache indirection needs to be updated - if beam_width > 1: - self._update_cache_indirection_buffer(scheduled_requests) - - decoder_input_buffers = self.store["decoder_input_buffers"][self.micro_batch_idx] # type: ignore - decoder_state = self.store["decoder_state"] - - make_decoding_batch_input( - decoder_input_buffers, - decoder_state, - scheduled_requests.context_requests, - scheduled_requests.generation_requests, - model_outputs["logits"], - beam_width, - num_context_logits_prefix_sum, - self.store["buffer_manager"], - ) - - self.algs.decoder.forward_async( # type: ignore - decoder_state, - self.store["decoder_input_buffers"][self.micro_batch_idx], # type: ignore - ) - - sampling_requests = ( - scheduled_requests.context_requests_last_chunk + scheduled_requests.generation_requests - ) - - finalize_events = {} - gathered_ids = None - if beam_width > 1: - finished_sum_device = decoder_state.finished_sum # type: ignore[attr-defined] - - for request in sampling_requests: - if request.is_context_init_state: - continue - if finished_sum_device[request.seq_slot] == beam_width: - finalize_events[request.request_id] = self._finalize_request(request, False) - elif request.streaming: - finalize_events[request.request_id] = self._finalize_request(request, True) - gathered_ids = self._copy_to_host(decoder_state.gathered_ids) # type: ignore[attr-defined] - new_output_tokens = self._copy_to_host(decoder_state.all_new_tokens) # type: ignore[attr-defined] - finished_sum = self._copy_to_host(decoder_state.finished_sum) # type: ignore[attr-defined] - finish_reasons = self._copy_to_host(decoder_state.finish_reasons) # type: ignore[attr-defined] - sequence_lengths = self._copy_to_host(decoder_state.sequence_lengths) # type: ignore[attr-defined] - - log_probs = None - cum_log_probs = None - if any(request.py_return_log_probs for request in sampling_requests): - log_probs = self._copy_to_host(decoder_state.log_probs) # type: ignore[attr-defined] - cum_log_probs = self._copy_to_host(decoder_state.cum_log_probs) # type: ignore[attr-defined] - - device = SampleStateTensors(new_tokens=decoder_state.all_new_tokens) # type: ignore[attr-defined] - - host = SampleStateTensorsHostTRTLLM( - new_tokens=new_output_tokens, - finished_sum=finished_sum, - finish_reasons=finish_reasons, - sequence_lengths=sequence_lengths, - log_probs=log_probs, - cum_log_probs=cum_log_probs, - gathered_ids=gathered_ids, - ) - - sampler_event = self._record_sampler_event() - - self.micro_batch_idx = (self.micro_batch_idx + 1) % self.num_micro_batches - - return SampleStateTRTLLM( - requests=sampling_requests, - device=device, - host=host, - sampler_event=sampler_event, - finalize_events=finalize_events, - ) - - @torch.inference_mode() - @override - def update_requests( - self, - state: SampleStateTRTLLM, - resource_manager: Optional[ResourceManager] = None, - ) -> None: - # resource_manager will not be used in this function, just for interface consistency. - assert isinstance(state, SampleStateTRTLLM) - - if state.sampler_event: - state.sampler_event.synchronize() - - if not state.requests: - return - - beam_width = self.beam_width(state.requests) - - if beam_width == 1 and self.MAX_DECODING_TOKENS == 1: - self.update_requests_single_beam_single_step(state) - else: - self.update_requests_multiple_beams_or_drafting(state, beam_width) - - @torch.inference_mode() - @nvtx_range("update_requests_single_beam_single_step") - def update_requests_single_beam_single_step(self, state: SampleStateTRTLLM) -> None: - """Specialization of update_requests for single beam and single step""" - assert state.host is not None - sequence_lengths_host_data = state.host.sequence_lengths.flatten().tolist() - finish_reasons = state.host.finish_reasons.flatten().tolist() - - reqs = [r for r in state.requests if not r.is_generation_complete_state] - - # NB: To ensure good performance, we must - # 1. Avoid accessing torch.Tensor object inside the for-each-request loops - # 2. Convert only necessary data to Python list - - # Add new tokens - reqs_with_new_tokens = [] - seq_slots = [] - seq_slots_need_log_probs = [] - for request in reqs: - assert request.py_seq_slot is not None - if sequence_lengths_host_data[request.py_seq_slot] <= request.get_num_tokens(0): - continue - - reqs_with_new_tokens.append(request) - seq_slots.append(request.py_seq_slot) - - if request.py_return_log_probs: - seq_slots_need_log_probs.append(request.py_seq_slot) - - # [maxTokensPerStep, batchSize, maxBeamWidth] - new_tokens = state.host.new_tokens[0, seq_slots, 0].tolist() - add_new_tokens_to_requests(reqs_with_new_tokens, new_tokens, 0) - - # Log probs - assert state.host is not None - if state.host.log_probs is not None: - # [batchSize, maxBeamWidth] - seq_last_idx = state.host.sequence_lengths[seq_slots_need_log_probs, 0] - 1 - # [batchSize, maxBeamWidth, maxSequenceLength] - log_probs_host = state.host.log_probs[ - seq_slots_need_log_probs, 0, seq_last_idx - ].tolist() - # [batchSize, maxBeamWidth] - assert state.host.cum_log_probs is not None - cum_log_probs_host = state.host.cum_log_probs[seq_slots_need_log_probs, 0].tolist() - - log_probs_idx = 0 - for request, new_token in zip(reqs_with_new_tokens, new_tokens): - if request.py_return_log_probs: - log_probs = [ - { - new_token: Logprob( - logprob=log_probs_host[log_probs_idx], - rank=1, - ) - } - ] - cum_log_probs = [cum_log_probs_host[log_probs_idx]] - request.py_result.append_log_probs([log_probs], cum_log_probs) - log_probs_idx += 1 - - for request in reqs: - request.py_decoding_iter += 1 - assert request.py_seq_slot is not None - finished_state = FinishedState(finish_reasons[request.py_seq_slot]) - if finished_state.is_finished: - request.state = LlmRequestState.GENERATION_COMPLETE - finish_reason = finished_state.to_finish_reason() - request.set_finished_reason(finish_reason, 0) - - @torch.inference_mode() - @nvtx_range("update_requests_multiple_beams_or_drafting") - def update_requests_multiple_beams_or_drafting( - self, - state: SampleStateTRTLLM, - beam_width: int, - ) -> None: - assert state.host is not None - new_tokens_host = state.host.new_tokens.tolist() - finished_sum_host = state.host.finished_sum.tolist() - finish_reasons = state.host.finish_reasons.flatten().tolist() - sequence_lengths_host_data = state.host.sequence_lengths.flatten().tolist() - cum_log_probs_host = ( - state.host.cum_log_probs.tolist() if state.host.cum_log_probs is not None else None - ) - log_probs_host = state.host.log_probs.tolist() if state.host.log_probs is not None else None - finalize_events = state.finalize_events - - reqs = [r for r in state.requests if not r.is_generation_complete_state] - - for request in reqs: - seq_slot = request.py_seq_slot - assert seq_slot is not None - num_generated_tokens = request.num_draft_tokens + 1 - current_num_of_tokens = request.max_beam_num_tokens - num_new_tokens = [0] * beam_width - - log_probs: list[list[dict[int, Logprob]]] = [[] for _ in range(beam_width)] - cum_log_probs = [] - - for beam_idx in range(beam_width): - seq_len = sequence_lengths_host_data[seq_slot * beam_width + beam_idx] - num_new_tokens[beam_idx] = min( - num_generated_tokens, seq_len - request.get_num_tokens(beam_idx) - ) - - for step in range(num_new_tokens[beam_idx]): - new_token = add_token(request, new_tokens_host, beam_idx=beam_idx, step=step) - - if request.py_return_log_probs: - assert state.host.log_probs is not None - assert log_probs_host is not None - # NOTE: Log probs with drafting has not been tested yet. - begin_log_probs_offset = ( - request.prompt_len if request.py_beam_width == 1 else 0 - ) - current_token = ( - seq_len - request.prompt_len - num_new_tokens[beam_idx] + step - ) - log_probs[beam_idx].append( - { - new_token: Logprob( - logprob=log_probs_host[seq_slot][beam_idx][ - begin_log_probs_offset + current_token - ], - rank=1, - ) - } - ) - - if request.py_return_log_probs: - assert cum_log_probs_host is not None - cum_log_probs.append(cum_log_probs_host[seq_slot][beam_idx]) - - finished_state = FinishedState(finish_reasons[seq_slot * beam_width + beam_idx]) - if finished_state.is_finished: - finish_reason = finished_state.to_finish_reason() - request.set_finished_reason(finish_reason, beam_idx) - - if request.py_return_log_probs: - request.py_result.append_log_probs(log_probs, cum_log_probs) - - # Set number of tokens predicted per runtime iteration. Will be > 1 for speculative decoding. - request.update_num_tokens_per_iteration( - request.max_beam_num_tokens - current_num_of_tokens, self.model_config - ) - - # Increment the decoding iteration counter - if request.state != LlmRequestState.GENERATION_COMPLETE: - request.py_decoding_iter += 1 - - if finished_sum_host[seq_slot] == beam_width: - request.state = LlmRequestState.GENERATION_COMPLETE - for request in reqs: - if finalize_events is not None and request.request_id in finalize_events: - self._post_process_request(request, state) - - def _finalize_request( - self, - request: LlmRequest, - streaming: bool, - ) -> CudaEvent: - """Finalizes the request. This is necessary for beam search.""" - seq_slot = request.py_seq_slot - event = cast( - CudaEvent, - self.algs.decoder.finalize( # type: ignore - self.store["decoder_state"], seq_slot, request.sampling_config, streaming - ), - ) - return event - - def _post_process_request(self, request: LlmRequest, state: SampleStateTRTLLM) -> None: - """Post Process the request. Updates the sequence according to the beam search results. - request: LlmRequest which shall be post processed - finalize_event: CudaEvent to wait for the finalize step to finish - """ - assert state.host is not None - seq_slot = request.py_seq_slot - beam_width = request.py_beam_width - # synchronize on the finalize event before continuing the post processing. - # should be unnecessary, as already wait for the sampler event in update_requests - assert state.finalize_events is not None - state.finalize_events[request.request_id].synchronize() - - # Get these values again, as they might have changed during the finalize step - output_ids_host = state.host.gathered_ids - assert output_ids_host is not None - sequence_lengths_host = state.host.sequence_lengths - - if request.py_return_log_probs: - log_probs_host = state.host.log_probs - cum_log_probs_host = state.host.cum_log_probs - else: - log_probs_host = None - cum_log_probs_host = None - - generated_tokens = [[0]] * beam_width - log_probs: list[list[dict[int, Logprob]]] = [[] for _ in range(beam_width)] - cum_log_probs = [] - - for beam_idx in range(beam_width): - # get the correct generated tokens for beam search - begin = request.py_prompt_len - end = cast(int, sequence_lengths_host[seq_slot, beam_idx].item()) - generated_tokens[beam_idx] = output_ids_host[seq_slot, beam_idx][begin:end].tolist() - - # get the correct log probs for beam search - if request.py_return_log_probs: - assert log_probs_host is not None - assert cum_log_probs_host is not None - cum_log_probs.append(cum_log_probs_host[seq_slot, beam_idx].item()) - - begin_log_probs_offset = request.prompt_len if request.py_beam_width == 1 else 0 - for current_token, token in enumerate(generated_tokens[beam_idx]): - log_probs[beam_idx].append( - { - token: Logprob( - logprob=log_probs_host[seq_slot, beam_idx][ - begin_log_probs_offset + current_token - ].item(), - rank=1, - ) - } - ) - if request.py_return_log_probs: - request.py_result.set_log_probs(log_probs, cum_log_probs) - - request.set_generated_tokens(generated_tokens) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py index 9200e0836039..d806a0e386bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py @@ -185,10 +185,11 @@ def request_random_seed(request: LlmRequest) -> Optional[int]: def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: sampling_config = request.sampling_config # These sampling fields live on the C++ SamplingConfig as optional> - # (a shape designed for the batched TRT-LLM sampler); the torch sampler consumes - # them per request, so we unwrap the singleton lists into scalars here. When the - # TRT-LLM sampler is removed, this SamplingConfig-based plumbing should be removed - # too in favor of reading the values directly from the per-request params. + # (a shape inherited from the batched C++ decoder that has since been + # removed); the torch sampler consumes them per request, so we unwrap the + # singleton lists into scalars here. + # TODO: drop this SamplingConfig-based plumbing in favor of reading the + # values directly from the per-request params. temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py index 296b95ad9717..6b203e005a88 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py @@ -24,9 +24,9 @@ deciding when a request is finished. * **Logit adjustments** -- embedding bias, d2t remapping and the fused greedy sampling kernel. -* **Async D2H** -- ``AsyncWorkerMixin`` (shared by ``TorchSampler`` and - ``TRTLLMSampler``), its private side-stream copier, and the ``SamplerEvent`` - that bundles the resulting worker futures / CUDA events for callers to await. +* **Async D2H** -- ``AsyncWorkerMixin`` (used by ``TorchSampler``), its + private side-stream copier, and the ``SamplerEvent`` that bundles the + resulting worker futures / CUDA events for callers to await. Anything here that outgrows a few dozen lines, or acquires per-slot state of its own, should move to its own module -- as beam search, penalties, token bans diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/seed_manager.py b/tensorrt_llm/_torch/pyexecutor/sampler/seed_manager.py index 67ce27997464..977e2c1aac0b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/seed_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/seed_manager.py @@ -55,8 +55,7 @@ class _SeedManager: is the first row of its strategy group. ``observe`` emits a one-time warning when a seeded request is seen. The per-row state is kept here so that honoring ``SamplingParams.seed`` becomes a FlashInfer version bump - rather than a redesign. ``TRTLLMSampler`` is unaffected -- its C++ - ``curandBatchInitialize`` seeds each slot's state individually. + rather than a redesign. Upstream fix in progress: https://github.com/flashinfer-ai/flashinfer/pull/2345 ("add per-request generator support for sampling kernels"), which also @@ -157,10 +156,9 @@ def observe(self, requests: list[LlmRequest]) -> None: "seed/offset per sampling call and distinguishes rows " "internally, so when several requests are sampled together " "only the first row's seed applies. Seeded requests are " - "therefore not yet reproducible unless sampled alone. Use " - "the TRTLLM sampler for fully per-request seeding; " - "TorchSampler support will land once FlashInfer honors " - "per-row seeds (tracked in " + "therefore not yet reproducible unless sampled alone. " + "Full per-request seeding will land once FlashInfer " + "honors per-row seeds (tracked in " "https://github.com/flashinfer-ai/flashinfer/pull/2345).", key="torch_sampler_per_request_seed_unsupported", ) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index c459d1faa1b3..a30fd832670d 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -204,8 +204,11 @@ def __init__(self, self.metrics_dict = {} self.candidate_metrics: list[dict] = [] self.trace_headers: Optional[dict[str, str]] = None - # torch backend will use trtllm sampler in beam search mode, but it does not support return logprobs incrementally - self.use_trtllm_sampler = sampling_params.use_beam_search and sampling_params.best_of > 1 + # Multi-beam search does not report logprobs incrementally: each response + # carries the full list, so it is sliced against _last_logprobs_len rather + # than appended wholesale. + self._logprobs_reported_cumulatively = (sampling_params.use_beam_search + and sampling_params.best_of > 1) if has_event_loop(): self.aqueue = AsyncQueue() @@ -295,8 +298,8 @@ def _maybe_fill_spec_dec_perf_metrics( drafting ran. The PyTorch executor instead attaches cumulative (accepted, drafted) totals to the response (LlmResult.spec_dec_totals, stashed on self in _handle_response); fill the section from them. - No-op when the section is already populated (TRT engine / TRTLLMSampler - paths) or when no drafting occurred. + No-op when the section is already populated (TRT engine path) or when + no drafting occurred. """ if not self.spec_dec_totals: return @@ -360,7 +363,7 @@ def _handle_sequence(self, *self._get_decoder_output_prefix_logprobs(), *response_tensors.log_probs[src_idx], ] - elif self.use_trtllm_sampler: + elif self._logprobs_reported_cumulatively: assert output._last_logprobs_len <= len( response_tensors.log_probs[src_idx] ), (f"_last_logprobs_len ({output._last_logprobs_len}) > log_probs length (" @@ -375,7 +378,7 @@ def _handle_sequence(self, # overcome some WAR in the cpp executor if finish_reasons[src_idx] != tllm.FinishReason.CANCELLED: - if self.use_trtllm_sampler and len( + if self._logprobs_reported_cumulatively and len( output.logprobs) > output.length: # LlmResult holds a reference to LogProbStorage, which may be updated by the worker before the result is serialized. # Therefore, we treat extra logprobs/logits as expected and only consume what's needed. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 645d715e721a..cf3312eb1441 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5133,13 +5133,6 @@ def validate_gms_config(self) -> 'GmsConfig': return self -class SamplerType(StrEnum): - """Enum for sampler type options.""" - TRTLLMSampler = "TRTLLMSampler" - TorchSampler = "TorchSampler" - auto = "auto" - - class PrefillCudaGraphBackend(StrEnum): """CUDA graph implementation used for prefill requests.""" @@ -5350,17 +5343,6 @@ def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': # tensorrt_llm/_torch/attention_backend/utils.py. telemetry=TelemetryField.categorical("VANILLA", "TRTLLM", "FLASHINFER")) - sampler_type: Union[str, SamplerType] = Field( - default=SamplerType.auto, - description= - "The type of sampler to use. Options are TRTLLMSampler, TorchSampler or auto. Defaults to auto, which will use TorchSampler. " - "TRTLLMSampler is deprecated and will be removed in release 1.4.", - status="deprecated", - deprecated= - "This parameter will be removed in release 1.4. TorchSampler will be the default sampler.", - telemetry=TelemetryField.categorical('TRTLLMSampler', 'TorchSampler', - 'auto')) - sampler_force_async_worker: bool = Field( default=False, description="Force usage of the async worker in the sampler for D2H " diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 0430e464a01d..60fd397f1149 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1454,17 +1454,6 @@ "kind": "value", "path": "sampler_force_async_worker" }, - { - "allowed_values": [ - "TRTLLMSampler", - "TorchSampler", - "auto" - ], - "annotation": "Union[str, tensorrt_llm.llmapi.llm_args.SamplerType]", - "converter": "allowlist", - "kind": "categorical", - "path": "sampler_type" - }, { "allowed_values": [ "MAX_UTILIZATION", diff --git a/tensorrt_llm/usage/schemas/README.md b/tensorrt_llm/usage/schemas/README.md index 01d406cdb4a6..df908c97e014 100644 --- a/tensorrt_llm/usage/schemas/README.md +++ b/tensorrt_llm/usage/schemas/README.md @@ -191,7 +191,6 @@ exhaustive field table at docs build time under **Developer Guide > Telemetry**. | `speculative_config.decoding_type` | Speculative decoding mode discriminator (e.g. `User_Provided`); other arms expose their own numeric/boolean knobs under `speculative_config.*`. | | `sparse_attention_config.algorithm` | Sparse attention algorithm discriminator; arm-specific knobs appear under `sparse_attention_config.*`. | | `reasoning_parser` | Reasoning parser selection, captured through an allowlist mirroring the `ReasoningParserFactory` registry. | -| `sampler_type` | Sampler selection, captured through an allowlist mirroring the `SamplerType` enum. | `llmApiConfigMetaJson` describes the capture process itself. It includes contract/version fields, schema and manifest digests, source args class, field diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 416f7f65463c..cd54eed25817 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -989,10 +989,8 @@ "test_e2e.py::test_openai_chat_harmony_perf_metrics": 161.03756944444444, "test_e2e.py::test_openai_chat_multimodal_example": 127.38494212962964, "test_e2e.py::test_openai_chat_with_logit_bias[torch_sampler]": 157.82800883002207, - "test_e2e.py::test_openai_chat_with_logit_bias[trtllm_sampler]": 157.81243107221007, "test_e2e.py::test_openai_completions_example[pytorch]": 311.7045913242009, "test_e2e.py::test_openai_completions_with_logit_bias[torch_sampler]": 70.47549118942732, - "test_e2e.py::test_openai_completions_with_logit_bias[trtllm_sampler]": 70.42805676855895, "test_e2e.py::test_openai_health": 60.098099999999995, "test_e2e.py::test_openai_kv_cache_contamination": 1219.6733020833333, "test_e2e.py::test_openai_lora": 171.0378888888889, @@ -1292,7 +1290,6 @@ "unittest/_torch/sampler/test_penalties.py": 61.55893835616438, "unittest/_torch/sampler/test_token_ban.py": 23.330488584474885, "unittest/_torch/sampler/test_torch_sampler.py": 251.46527294117647, - "unittest/_torch/sampler/test_trtllm_sampler.py": 78.89402733485194, "unittest/_torch/speculative/hw_agnostic": 192.419382629108, "unittest/_torch/speculative/test_capture_override_leak.py": 18.178845569620254, "unittest/_torch/speculative/test_dspark_cute_dsl_attention.py": 27.91303125, diff --git a/tests/integration/defs/ray_orchestrator/RL/run_rl_perf_reproduce.py b/tests/integration/defs/ray_orchestrator/RL/run_rl_perf_reproduce.py index c0d17056e5af..0567656a6898 100644 --- a/tests/integration/defs/ray_orchestrator/RL/run_rl_perf_reproduce.py +++ b/tests/integration/defs/ray_orchestrator/RL/run_rl_perf_reproduce.py @@ -58,7 +58,6 @@ async def init_llm(self): tensor_parallel_size=self.async_llm_kwargs["tensor_parallel_size"], trust_remote_code=self.async_llm_kwargs["trust_remote_code"], sleep_config=self.async_llm_kwargs["sleep_config"], - sampler_type=self.async_llm_kwargs["sampler_type"], placement_groups=self.async_llm_kwargs["placement_groups"], placement_bundle_indices=self.async_llm_kwargs["placement_bundle_indices"], per_worker_gpu_share=self.async_llm_kwargs["per_worker_gpu_share"], @@ -201,7 +200,6 @@ async def setup_rl_llm(args): ExecutorMemoryType.KV_CACHE: "NONE", } ), - "sampler_type": args.sampler_type, "placement_groups": placement_group_list[i], "placement_bundle_indices": placement_bundle_indices_list[i], "per_worker_gpu_share": 0.5, @@ -277,13 +275,6 @@ def add_rl_llm_args(parser): parser.add_argument( "--max_num_tokens", type=int, default=32768, help="Maximum number of tokens." ) - parser.add_argument( - "--sampler_type", - type=str, - default="TRTLLMSampler", - choices=["TRTLLMSampler", "TorchSampler"], - help="Sampler type.", - ) parser.add_argument( "--trust_remote_code", action="store_true", diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 32ef7468697f..b84ccbac1fe0 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -693,21 +693,20 @@ def test_openai_post_processor(llm_root, llm_venv): str(test_root / "_test_openai_post_processor.py")]) -@pytest.mark.parametrize("sampler", ["torch_sampler", "trtllm_sampler"]) -def test_openai_completions_with_logit_bias(llm_root, llm_venv, sampler: str): +def test_openai_completions_with_logit_bias(llm_root, llm_venv): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ "-m", "pytest", - str(test_root / "_test_openai_completions.py"), "-k", sampler + str(test_root / "_test_openai_completions.py"), "-k", + "logit_bias_effect" ]) -@pytest.mark.parametrize("sampler", ["torch_sampler", "trtllm_sampler"]) -def test_openai_chat_with_logit_bias(llm_root, llm_venv, sampler: str): +def test_openai_chat_with_logit_bias(llm_root, llm_venv): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ "-m", "pytest", - str(test_root / "_test_openai_chat.py"), "-k", sampler + str(test_root / "_test_openai_chat.py"), "-k", "logit_bias_effect" ]) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 2488cec777ed..df5ffd99bea0 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -26,7 +26,6 @@ l0_a10: - unittest/_torch/modeling/test_modeling_radio.py - unittest/_torch/modeling/test_modeling_minicpmv4_6.py - unittest/_torch/modeling/test_multimodal_encoder_mixin.py - - unittest/_torch/sampler/test_trtllm_sampler.py - unittest/_torch/sampler/test_token_ban.py - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/test_kv_cache_compression_manager.py diff --git a/tests/integration/test_lists/test-db/l0_a30.yml b/tests/integration/test_lists/test-db/l0_a30.yml index 3a5374fa3592..129f3d14819d 100644 --- a/tests/integration/test_lists/test-db/l0_a30.yml +++ b/tests/integration/test_lists/test-db/l0_a30.yml @@ -22,10 +22,8 @@ l0_a30: - unittest/_torch/sampler/test_beam_search.py - unittest/_torch/sampler/test_beam_search_speculative_d2h.py - unittest/_torch/sampler/test_logits_logprobs.py - - test_e2e.py::test_openai_completions_with_logit_bias[torch_sampler] - - test_e2e.py::test_openai_chat_with_logit_bias[torch_sampler] - - test_e2e.py::test_openai_completions_with_logit_bias[trtllm_sampler] - - test_e2e.py::test_openai_chat_with_logit_bias[trtllm_sampler] + - test_e2e.py::test_openai_completions_with_logit_bias + - test_e2e.py::test_openai_chat_with_logit_bias - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - condition: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 55d61f013189..f8eecd886bc8 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -219,10 +219,8 @@ full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_ full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestGemma3_27BInstruct::test_fp8_prequantized SKIP (https://nvbugs/6479708) full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6657522) -full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[block_reuse-python_scheduler-TRTLLMSampler] SKIP (https://nvbugs/6608387) -full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[block_reuse-python_scheduler-TorchSampler] SKIP (https://nvbugs/6608387) -full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[no_reuse-cpp_scheduler-TRTLLMSampler] SKIP (https://nvbugs/6608387) -full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[no_reuse-python_scheduler-TorchSampler] SKIP (https://nvbugs/6608387) +full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[block_reuse-python_scheduler] SKIP (https://nvbugs/6608387) +full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[no_reuse-python_scheduler] SKIP (https://nvbugs/6608387) full:GB300/unittest/_torch/modeling/test_modeling_gpt_oss.py::test_gpt_oss_trtllmgen[CUTLASS] SKIP (https://nvbugs/6633932) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6313072) @@ -349,15 +347,13 @@ test_e2e.py::test_trtllm_bench_llmapi_launch[pytorch_backend-llama-v3-llama3-8b] unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6641268) -unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_block_reuse_cache_hit[TorchSampler] SKIP (https://nvbugs/6608387) +unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_block_reuse_cache_hit SKIP (https://nvbugs/6608387) unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy SKIP (https://nvbugs/6607482) unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py::test_nemotron_nano_v2_vl_video_batch_equivalence SKIP (https://nvbugs/6625695) unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) unittest/_torch/modules/test_w4a16_nvfp4_linear.py::test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin SKIP (https://nvbugs/6581071) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404) -unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TRTLLMSampler-cuda_graph_and_overlap-None-1-1-True-True-False] SKIP (https://nvbugs/6463819) -unittest/_torch/sampler/test_trtllm_sampler.py::test_trtllm_sampler_best_of_with_logprobs SKIP (https://nvbugs/6487837) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cublaslt] SKIP (https://nvbugs/6581067) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cutlass] SKIP (https://nvbugs/6581067) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cublaslt] SKIP (https://nvbugs/6581067) diff --git a/tests/unittest/_torch/executor/test_overlap_scheduler.py b/tests/unittest/_torch/executor/test_overlap_scheduler.py index 08bbb3a087a8..7f9a57f6560c 100644 --- a/tests/unittest/_torch/executor/test_overlap_scheduler.py +++ b/tests/unittest/_torch/executor/test_overlap_scheduler.py @@ -24,14 +24,12 @@ def model_path(): def create_llm(model_dir, disable_overlap_scheduler, - sampler_type, scheduler_config=None, enable_block_reuse=False): """Create LLM with specific overlap scheduler setting""" if scheduler_config is None: scheduler_config = SchedulerConfig() - pytorch_config = dict(disable_overlap_scheduler=disable_overlap_scheduler, - sampler_type=sampler_type) + pytorch_config = dict(disable_overlap_scheduler=disable_overlap_scheduler) trt_kv_cache_config = TRT_KvCacheConfig( enable_block_reuse=enable_block_reuse) @@ -50,14 +48,13 @@ def create_llm(model_dir, ) -@pytest.mark.parametrize("sampler_type", ["TorchSampler", "TRTLLMSampler"]) @pytest.mark.parametrize("use_python_scheduler", [False, True], ids=["cpp_scheduler", "python_scheduler"]) @pytest.mark.parametrize("enable_block_reuse", [False, True], ids=["no_reuse", "block_reuse"]) @pytest.mark.high_cuda_memory @pytest.mark.mpi_ray_parity -def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, +def test_overlap_scheduler_consistency(model_path, test_case, use_python_scheduler, enable_block_reuse): scheduler_config = SchedulerConfig( @@ -80,7 +77,6 @@ def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, # Test with overlap scheduler enabled with create_llm(model_path, disable_overlap_scheduler=False, - sampler_type=sampler_type, scheduler_config=scheduler_config, enable_block_reuse=enable_block_reuse) as llm: outputs_with_overlap = llm.generate(prompts, @@ -93,7 +89,6 @@ def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, # Test with overlap scheduler disabled with create_llm(model_path, disable_overlap_scheduler=True, - sampler_type=sampler_type, scheduler_config=scheduler_config, enable_block_reuse=enable_block_reuse) as llm: outputs_without_overlap = llm.generate(prompts, @@ -110,11 +105,9 @@ def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, assert with_overlap == without_overlap -@pytest.mark.parametrize("sampler_type", ["TorchSampler", "TRTLLMSampler"]) @pytest.mark.high_cuda_memory @pytest.mark.mpi_ray_parity -def test_overlap_scheduler_block_reuse_cache_hit(model_path, test_case, - sampler_type): +def test_overlap_scheduler_block_reuse_cache_hit(model_path, test_case): """Verify that blocks are actually reused when sending the same prompt twice with the overlap scheduler enabled. Uses a single prompt to avoid batch-internal cache hits that could make the cold-cache check flaky.""" @@ -133,7 +126,6 @@ def test_overlap_scheduler_block_reuse_cache_hit(model_path, test_case, with create_llm(model_path, disable_overlap_scheduler=False, - sampler_type=sampler_type, enable_block_reuse=True) as llm: output_first = llm.generate([prompt], sampling_params=sampling_config, diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index 67ff962fe099..b69ced59f4c5 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -200,7 +200,6 @@ def _make_llm_args(): calibration_file_path=None, calibration_layer_indices=None, ), - sampler_type=None, cuda_graph_config=None, parallel_config=SimpleNamespace(to_mapping=lambda: SimpleNamespace()), get_runtime_sizes=lambda: (1, 128, 128, 4), diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index af33d59bbb9c..6b6d0b5a7019 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -63,13 +63,8 @@ def fixed_params(): return {"max_tokens": 8, "max_beam_width": 2} -@pytest.fixture(scope="module", params=["TRTLLMSampler", "TorchSampler"]) -def sampling_information(request): - return request.param - - @pytest.fixture(scope="module") -def model_kwargs(fixed_params, sampling_information) -> dict[str, Any]: +def model_kwargs(fixed_params) -> dict[str, Any]: assert fixed_params[ "max_beam_width"] == 2, "This test only works for a beam width of 2" @@ -79,7 +74,6 @@ def model_kwargs(fixed_params, sampling_information) -> dict[str, Any]: weight_loader=DummyWeightLoader(), config_loader=DummyConfigLoader(), ), - sampler_type=sampling_information, ) @@ -130,8 +124,6 @@ def _single_process_context(): def llm(fixed_params, input_prompts, model_kwargs, single_process: bool, with_cuda_graph_and_overlap: bool): check_no_sync = single_process # single_process only used for sync check - if check_no_sync and model_kwargs["sampler_type"] != "TorchSampler": - pytest.skip("Sync check only supported for TorchSampler") gc.collect( 2) # force destruction of any other LLM instances (cf. comment above) @@ -405,19 +397,6 @@ def test_beam_search_e2e( ) -> None: llm_args = cast(TorchLlmArgs, llm.args) # type: ignore[redundant-cast] - if return_log_probs and num_prompts > 1 and llm_args.sampler_type == "TRTLLMSampler": - pytest.skip( - "Beam search currently does not support return_log_probs with multiple prompts" - ) - if return_log_probs and llm_args.sampler_type == "TRTLLMSampler": - pytest.skip( - "Beam search on TRTLLMSampler does not correctly handle log_probs if called multiple times" - ) - if stop_token_ids is not None and llm_args.sampler_type == "TRTLLMSampler": - pytest.skip( - "Beam search on TRTLLMSampler does not correctly handle stop_token_ids" - ) - # create sampling parameters # additional_model_outputs is used to gather the cache indirection from the model. sampling_params = SamplingParams( @@ -535,11 +514,6 @@ def test_beam_search_disagg_first_token_is_end_id( once to see what the beams sample, then declare beam 0's token the end id and rerun, so the context step finishes on its first and only token. """ - if model_kwargs["sampler_type"] != "TorchSampler": - pytest.skip( - "The context-side end-id mask is a TorchSampler path; the C++ " - "decoder behind TRTLLMSampler pools the end candidate instead.") - beam_width = fixed_params["max_beam_width"] base_params = SamplingParams( max_tokens=fixed_params["max_tokens"], @@ -640,7 +614,6 @@ def test_beam_search_large_beam_width_regression( llm = LLM( model=_pl.Path("dummy_path"), checkpoint_loader=checkpoint_loader, - sampler_type="TRTLLMSampler", max_beam_width=beam_width, max_batch_size=beam_width * num_prompts, max_seq_len=64, @@ -805,7 +778,6 @@ def recording_update_requests(self: TorchSampler, state, *args, **kwargs): llm = LLM( model=_pl.Path("dummy_path"), checkpoint_loader=checkpoint_loader, - sampler_type="TorchSampler", max_beam_width=max_beam_width, max_batch_size=max_beam_width, max_seq_len=64, @@ -1814,9 +1786,9 @@ def test_vbws_cpp_formula_matches_past_array_end(): kMaxBeamWidthArrayLength rather than the actual array length, so it read out of bounds and returned arbitrary widths (observed: 0, 32, 849 for a 3-entry array). That starved the request in the C++ micro-batch scheduler - and hung decoding; it is fixed in llmRequest.cpp. TRTLLMSampler and the - scheduler call into C++ directly, so pin the agreement here -- a failure - means the two clamps have drifted apart again. + and hung decoding; it is fixed in llmRequest.cpp. The scheduler calls into + C++ directly, so pin the agreement here -- a failure means the two clamps + have drifted apart again. """ beam_width_array = [2, 3, 4] request = _vbws_request(beam_width_array) @@ -2279,11 +2251,6 @@ def fixed_params(): def batch_size(request) -> int: return cast(int, request.param) - @pytest.fixture(scope="module", params=["TRTLLMSampler", "TorchSampler"]) - @staticmethod - def sampler_type(request) -> str: - return cast(str, request.param) - @pytest.fixture(scope="module") @staticmethod def model_kwargs() -> dict[str, Any]: @@ -2295,16 +2262,11 @@ def model_kwargs() -> dict[str, Any]: # NB: Class-level fixture overrides do not work without this @pytest.fixture(scope="module") @staticmethod - def llm(fixed_params, input_prompts, model_kwargs, batch_size: int, - sampler_type: str): + def llm(fixed_params, input_prompts, model_kwargs, batch_size: int): return _build_llm( fixed_params, input_prompts, - (model_kwargs - | dict( - max_batch_size=batch_size, - sampler_type=sampler_type, - )), + (model_kwargs | dict(max_batch_size=batch_size)), ) def _check_engine_responds(self, llm: LLM, input_prompts: list[str], @@ -2327,15 +2289,12 @@ def test_use_beam_search_disabled_rejects_multiple_returns( input_prompts: list[str], fixed_params: dict[str, Any], batch_size: int, - sampler_type: str, use_beam_search: bool | None, ): # best_of > 1 without beam search is greedy multi-return, which the LLM # API rejects. Covers use_beam_search both explicitly False and omitted. if batch_size == 1: pytest.skip("Test does not depend on batch size") - if sampler_type == "TorchSampler": - pytest.skip("Test does not depend on sampler_type") assert fixed_params["max_beam_width"] > 2 params = dict( max_tokens=fixed_params["max_tokens"], @@ -2363,7 +2322,6 @@ def test_exhaustive_early_stopping_allowed_without_disagg( input_prompts: list[str], fixed_params: dict[str, Any], batch_size: int, - sampler_type: str, early_stopping: int, ): # Beam search is rejected wholesale under disaggregated serving (the @@ -2373,8 +2331,6 @@ def test_exhaustive_early_stopping_allowed_without_disagg( # testing a bound method rather than calling it. if batch_size == 1: pytest.skip("Test does not depend on batch size") - if sampler_type == "TRTLLMSampler": - pytest.skip("Exhaustive early_stopping check is TorchSampler-side") outputs = llm.generate(input_prompts, sampling_params=SamplingParams( max_tokens=fixed_params["max_tokens"], @@ -2400,7 +2356,6 @@ def test_smaller_beam_width( input_prompts: list[str], fixed_params: dict[str, Any], batch_size: int, - sampler_type: str, ): if batch_size == 1: pytest.skip("Test does not depend on batch size") @@ -2437,35 +2392,6 @@ def test_smaller_beam_width( )) self._check_engine_responds(llm, input_prompts, fixed_params) - @pytest.mark.timeout(120) - @pytest.mark.threadleak(enabled=False) - def test_logprobs_trtllm_sampler( - self, - llm: LLM, - input_prompts: list[str], - fixed_params: dict[str, Any], - batch_size: int, - sampler_type: str, - ): - if sampler_type != "TRTLLMSampler": - pytest.skip("Test is specific to TRTLLMSampler") - - with pytest.raises( - RequestError, - match= - ".*Beam search only supports logprobs when batch size is 1.*" - ) if batch_size > 1 else nullcontext(): - _ = llm.generate(input_prompts, - sampling_params=SamplingParams( - max_tokens=fixed_params["max_tokens"], - n=1, - best_of=fixed_params["max_beam_width"], - use_beam_search=True, - end_id=-1, - logprobs=1, - )) - self._check_engine_responds(llm, input_prompts, fixed_params) - @pytest.mark.timeout(120) @pytest.mark.threadleak(enabled=False) def test_logprobs_torch_sampler( @@ -2474,10 +2400,7 @@ def test_logprobs_torch_sampler( input_prompts: list[str], fixed_params: dict[str, Any], batch_size: int, - sampler_type: str, ): - if sampler_type != "TorchSampler": - pytest.skip("Test is specific to TorchSampler") if batch_size == 1: pytest.skip("Test does not depend on batch size") diff --git a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py index fd3a38b53257..1b4f31e47c6a 100644 --- a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py +++ b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py @@ -72,7 +72,6 @@ def _build_llm( weight_loader=DummyWeightLoader(), config_loader=DummyConfigLoader(), ), - sampler_type="TorchSampler", max_batch_size=fixed_params["max_beam_width"] * len(input_prompts), kv_cache_config=KvCacheConfig(max_tokens=10000), # pyright: ignore max_seq_len=32, diff --git a/tests/unittest/_torch/sampler/test_logits_logprobs.py b/tests/unittest/_torch/sampler/test_logits_logprobs.py index 9637d8b47c8c..6c9d8048c288 100644 --- a/tests/unittest/_torch/sampler/test_logits_logprobs.py +++ b/tests/unittest/_torch/sampler/test_logits_logprobs.py @@ -59,11 +59,6 @@ def disable_overlap_scheduler_fixture(request) -> bool: return request.param -@pytest.fixture(scope="module", params=["TRTLLMSampler", "TorchSampler"]) -def sampler_type_fixture(request) -> str: - return request.param - - @pytest.fixture(scope="module", params=[False, True]) def enable_early_first_token_response_fixture(request) -> bool: return request.param @@ -92,11 +87,9 @@ def get_salt(cls, reuse_cache: bool) -> str: @pytest.fixture(scope="module") def llm( - sampler_type_fixture: str, disable_overlap_scheduler_fixture: bool, enable_early_first_token_response_fixture: bool, ): - sampler_type = sampler_type_fixture disable_overlap_scheduler = disable_overlap_scheduler_fixture enable_early_first_token_response = enable_early_first_token_response_fixture @@ -109,7 +102,6 @@ def llm( model=os.path.join(llm_models_root(), "llama-models-v2", "TinyLlama-1.1B-Chat-v1.0"), kv_cache_config=global_kvcache_config, max_batch_size=128, # reduce buffer sizes, specially for generation logits - sampler_type=sampler_type, disable_overlap_scheduler=disable_overlap_scheduler, enable_early_first_token_response=enable_early_first_token_response, ) diff --git a/tests/unittest/_torch/sampler/test_penalties_e2e.py b/tests/unittest/_torch/sampler/test_penalties_e2e.py index 553aa51717c4..05222a3d64f8 100644 --- a/tests/unittest/_torch/sampler/test_penalties_e2e.py +++ b/tests/unittest/_torch/sampler/test_penalties_e2e.py @@ -125,7 +125,6 @@ def _create_torch_llm( trust_remote_code=True, enable_chunked_prefill=True, cuda_graph_config=CudaGraphConfig(), - sampler_type="TorchSampler", kv_cache_config=TRT_KvCacheConfig(enable_block_reuse=False), max_num_tokens=128, enable_iter_perf_stats=enable_iter_perf_stats, @@ -412,7 +411,6 @@ def test_beam_search_penalties_e2e(overlap: bool) -> None: weight_loader=DummyWeightLoader(), config_loader=DummyConfigLoader(), ), - sampler_type="TorchSampler", max_batch_size=_E2E_BEAM_WIDTH, kv_cache_config=TRT_KvCacheConfig(max_tokens=10000), max_seq_len=32, diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 28394f70425a..cd4cd7750642 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -3422,9 +3422,8 @@ class TestTopPDecay: """Minimal functional guards for Top-P Decay in TorchSampler. Covers strategy routing, the post-sample runtime update (parity with the - C++ computeToppDecay recurrence; cases ported from - topPSamplingLayerTest.cpp), and per-request rejection of unsupported - combinations. + recurrence the former C++ computeToppDecay implemented), and per-request + rejection of unsupported combinations. """ VOCAB_SIZE = 1000 diff --git a/tests/unittest/_torch/sampler/test_trtllm_sampler.py b/tests/unittest/_torch/sampler/test_trtllm_sampler.py deleted file mode 100644 index 032a7bc21697..000000000000 --- a/tests/unittest/_torch/sampler/test_trtllm_sampler.py +++ /dev/null @@ -1,199 +0,0 @@ -import pytest -from utils.llm_data import llm_models_root -from utils.util import similar - -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.llmapi import CudaGraphConfig -from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig - - -@pytest.fixture(scope="module") -def model_path(): - return llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" - - -def _create_llm_base(model_dir, enable_trtllm_sampler): - """Base LLM creation with configurable sampler.""" - sampler_type = "TRTLLMSampler" if enable_trtllm_sampler else "TorchSampler" - - trt_kv_cache_config = TRT_KvCacheConfig(enable_block_reuse=False) - - return LLM( - model=str(model_dir), - tensor_parallel_size=1, - trust_remote_code=True, - enable_chunked_prefill=True, - cuda_graph_config=CudaGraphConfig(), - sampler_type=sampler_type, - kv_cache_config=trt_kv_cache_config, - max_num_tokens= - 128 # Only one request longer than max_num_tokens is required to test chunked prefill - ) - - -def create_llm(model_dir): - """Create LLM with specific overlap scheduler setting""" - return _create_llm_base(model_dir, enable_trtllm_sampler=True) - - -def create_llm_with_torch_sampler(model_dir): - """Create LLM with TorchSampler.""" - return _create_llm_base(model_dir, enable_trtllm_sampler=False) - - -@pytest.mark.high_cuda_memory -def test_trtllm_sampler(model_path): - prompts = [ - "Magellan and Elcano lead the first", - "The capital of France is", - "The capital of Bolivia is", - ] - - expected_outputs = [["circumnavigation of the world"], ["Paris"], - ["La Paz"]] - - # Test configuration - max_new_tokens = 10 - temperature = 1.0 - top_p = None - stop_words = ["."] - - sampling_config = SamplingParams(max_tokens=max_new_tokens, - n=1, - stop=stop_words, - temperature=temperature, - top_p=top_p) - - # Test with overlap scheduler disabled - llm = create_llm(model_path) - outputs = llm.generate(prompts, - sampling_params=sampling_config, - use_tqdm=True) - texts = [[completion.text for completion in request_output.outputs] - for request_output in outputs] - llm.shutdown() - - # Remove any text after \n\n, consider texts is a list of list of strings - texts = [[text.split('\n\n')[0] for text in request_output] - for request_output in texts] - - # Verify outputs are consistent - for text, expected in zip(texts, expected_outputs): - assert similar(text, expected), f"text: {text}, expected: {expected}" - - -@pytest.mark.high_cuda_memory -def test_trtllm_sampler_with_stop_token_ids(model_path): - """Test sampler with stop_token_ids (fast path optimization).""" - - llm = create_llm_with_torch_sampler(model_path) - tokenizer = llm.tokenizer - - prompt = "The capital of France is" - target_sentence = "The capital of France is Paris" - - prompt_tokens = tokenizer.encode(prompt, add_special_tokens=False) - target_tokens = tokenizer.encode(target_sentence, add_special_tokens=False) - - # Use the first token after the prompt as the stop token - assert len(target_tokens) > len( - prompt_tokens), "Target must be longer than prompt" - stop_token_id = target_tokens[len(prompt_tokens)] - - sampling_config = SamplingParams(max_tokens=100, - n=1, - stop_token_ids=[stop_token_id], - temperature=0.0) - - outputs = llm.generate([prompt], sampling_params=sampling_config) - text = outputs[0].outputs[0].text - - output_tokens = tokenizer.encode(text, add_special_tokens=False) - - llm.shutdown() - assert stop_token_id not in output_tokens, f"Output should not contain stop token {stop_token_id}" - assert len(output_tokens - ) < 10, "Should stop very early with first-token stop_token_id" - - -@pytest.mark.high_cuda_memory -def test_torch_sampler_with_multi_token_stop_words(model_path): - """Test TorchSampler with multi-token stop words (slow path).""" - - llm = create_llm_with_torch_sampler(model_path) - tokenizer = llm.tokenizer - - prompt = "The capital of France is" - - # Use a string that will tokenize to multiple tokens - stop_string = "\n\n" - stop_tokens = tokenizer.encode(stop_string, add_special_tokens=False) - - assert len( - stop_tokens - ) > 1, f"Stop string should be multi-token, got {len(stop_tokens)} tokens" - - sampling_config = SamplingParams( - max_tokens=100, - n=1, - stop=[stop_string], # Use 'stop' parameter for multi-token - temperature=0.0) - - outputs = llm.generate([prompt], sampling_params=sampling_config) - text = outputs[0].outputs[0].text - - llm.shutdown() - - assert len(text) > 0, "Should generate some text" - assert stop_string not in text, f"Stop string '{repr(stop_string)}' should not appear in the output" - - -@pytest.mark.high_cuda_memory -def test_trtllm_sampler_best_of_with_logprobs(model_path): - """Test TRTLLMSampler with best_of > n and logprobs.""" - - llm = create_llm(model_path) - - prompt = "The capital of France is" - - sampling_config = SamplingParams( - max_tokens=10, - temperature=1.0, - top_k=2, - n=2, # Return 2 sequences - best_of=3, # Generate 3 candidates, pick best 2 - logprobs=1 # Return log probabilities - ) - - outputs = llm.generate([prompt], sampling_params=sampling_config) - - llm.shutdown() - - assert len(outputs) == 1, "Should return one request output" - - request_output = outputs[0] - completion_outputs = request_output.outputs - - assert len( - completion_outputs - ) == 2, f"Expected 2 outputs (n=2), got {len(completion_outputs)}" - - for i, output in enumerate(completion_outputs): - assert len(output.text) > 0, f"Output {i} should have generated text" - - assert output.finish_reason is not None, \ - f"Output {i} must have a finish_reason" - - assert output.cumulative_logprob is not None, \ - f"Output {i} should have cumulative_logprob when logprobs is requested" - assert isinstance(output.cumulative_logprob, (float, int)), \ - f"Output {i} cumulative_logprob should be a number, got {type(output.cumulative_logprob)}" - - assert output.logprobs is not None, \ - f"Output {i} should have logprobs when logprobs=1" - assert len(output.logprobs) == len(output.token_ids), \ - f"Output {i} should have logprobs for each token" - - if len(completion_outputs) >= 2: - assert completion_outputs[0].cumulative_logprob >= completion_outputs[1].cumulative_logprob, \ - "Outputs should be sorted by cumulative log probability (best first)" diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index 30f3727f688f..66fd496635b1 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -32,7 +32,7 @@ from tensorrt_llm.llmapi import (CalibConfig, CompletionOutput, GuidedDecodingParams, QuantConfig, RequestOutput, SamplingParams) -from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend, SamplerType +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend from tensorrt_llm.llmapi.llm_utils import LlmArgs from tensorrt_llm.logger import Singleton from tensorrt_llm.sampling_params import LogprobMode diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d66b9e72dafe..d74aa6054cd9 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -171,10 +171,6 @@ methods: annotation: str default: TRTLLM status: beta - sampler_type: - annotation: Union[str, tensorrt_llm.llmapi.llm_args.SamplerType] - default: auto - status: deprecated sampler_force_async_worker: annotation: bool default: False diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_sampler.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_sampler.py deleted file mode 100644 index 1ba957d4cc9c..000000000000 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_sampler.py +++ /dev/null @@ -1,53 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from _model_test_utils import get_small_model_config -from build_and_run_ad import ExperimentConfig, main - -from tensorrt_llm.llmapi.llm_args import SamplerType - - -def test_ad_trtllm_sampler_smoke(): - """Test TRTLLMSampler in AutoDeploy smoke test.""" - # Get small model config - model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct" - experiment_config = get_small_model_config(model_id) - - # Configure for TRTLLMSampler - experiment_config["args"]["runtime"] = "trtllm" - experiment_config["args"]["world_size"] = 1 - # NOTE: trtllm attention backend fails on B200 (likely illegal memory access); use flashinfer. - experiment_config["args"]["attn_backend"] = "flashinfer" - experiment_config["args"]["sampler_type"] = SamplerType.TRTLLMSampler - - # Setup simple prompt - experiment_config["prompt"]["batch_size"] = 1 - experiment_config["prompt"]["queries"] = "What is the capital of France?" - experiment_config["prompt"]["sp_kwargs"] = { - "max_tokens": 10, - "temperature": 1.0, - "top_k": 1, - } - - print(f"Experiment config: {experiment_config}") - cfg = ExperimentConfig(**experiment_config) - - print("Running smoke test with TRTLLMSampler...") - results = main(cfg) - - # Basic assertion that we got some output - prompts_and_outputs = results["prompts_and_outputs"] - assert len(prompts_and_outputs) == 1 - assert len(prompts_and_outputs[0][1]) > 0 diff --git a/tests/unittest/auto_deploy/standalone/test_standalone_test_export.py b/tests/unittest/auto_deploy/standalone/test_standalone_test_export.py index 49525ff30cfd..676f4562b550 100644 --- a/tests/unittest/auto_deploy/standalone/test_standalone_test_export.py +++ b/tests/unittest/auto_deploy/standalone/test_standalone_test_export.py @@ -47,7 +47,6 @@ "test_ad_build_small_single.py": "test_paragraf_trtllm_build_small_single.py", "test_ad_guided_decoding_regex.py": "test_paragraf_trtllm_guided_decoding_regex.py", "test_ad_trtllm_bench.py": "test_paragraf_trtllm_bench.py", - "test_ad_trtllm_sampler.py": "test_paragraf_trtllm_sampler.py", "test_ad_trtllm_serve.py": "test_paragraf_trtllm_serve.py", } diff --git a/tests/unittest/executor/test_spec_dec_perf_metrics.py b/tests/unittest/executor/test_spec_dec_perf_metrics.py index 6a7d1cb6de13..b71539765976 100644 --- a/tests/unittest/executor/test_spec_dec_perf_metrics.py +++ b/tests/unittest/executor/test_spec_dec_perf_metrics.py @@ -65,7 +65,7 @@ def test_noop_without_totals(): def test_noop_when_section_already_populated(): - # TRT-engine / TRTLLMSampler paths populate the section runtime-side + # The TRT-engine path populates the section runtime-side # (updateNumTokensPerIteration); the backfill must not overwrite it. pm = tllm.RequestPerfMetrics() spec_dec = tllm.SpeculativeDecodingMetrics() diff --git a/tests/unittest/llmapi/apps/_test_openai_chat.py b/tests/unittest/llmapi/apps/_test_openai_chat.py index 78b52d785e96..db3b092c4adc 100644 --- a/tests/unittest/llmapi/apps/_test_openai_chat.py +++ b/tests/unittest/llmapi/apps/_test_openai_chat.py @@ -516,21 +516,9 @@ def test_stop_reason(client: openai.OpenAI, model_name: str, backend: str): @pytest.mark.asyncio(loop_scope='function') -@pytest.mark.parametrize( - 'server_with_custom_sampler', - [ - { - 'sampler_type': "TorchSampler" - }, # torch_sampler - { - 'sampler_type': "TRTLLMSampler" - }, # trtllm_sampler - ], - indirect=True, - ids=['torch_sampler', 'trtllm_sampler']) async def test_chat_completion_with_logit_bias_effect( server_with_custom_sampler, model_name: str) -> None: - '''Test that logit bias affects output as expected for both samplers (chat endpoint).''' + '''Test that logit bias affects output as expected (chat endpoint).''' client = server_with_custom_sampler.get_async_client() await logit_bias_effect_helper(client, model_name, 'chat') diff --git a/tests/unittest/llmapi/apps/_test_openai_completions.py b/tests/unittest/llmapi/apps/_test_openai_completions.py index 3fe8a6dd70bd..59bb56ea4bc5 100644 --- a/tests/unittest/llmapi/apps/_test_openai_completions.py +++ b/tests/unittest/llmapi/apps/_test_openai_completions.py @@ -407,22 +407,10 @@ async def test_completion_streaming(async_client: openai.AsyncOpenAI, @pytest.mark.asyncio(loop_scope='function') -@pytest.mark.parametrize( - 'server_with_custom_sampler', - [ - { - 'sampler_type': "TorchSampler" - }, # torch_sampler - { - 'sampler_type': "TRTLLMSampler" - }, # trtllm_sampler - ], - indirect=True, - ids=['torch_sampler', 'trtllm_sampler']) async def test_completion_with_logit_bias_effect( server_with_custom_sampler: RemoteOpenAIServer, model_name: str) -> None: - '''Test that logit bias affects output as expected for both samplers (completions endpoint).''' + '''Test that logit bias affects output as expected (completions endpoint).''' client = server_with_custom_sampler.get_async_client() await logit_bias_effect_helper(client, model_name, 'completions') diff --git a/tests/unittest/llmapi/apps/utils.py b/tests/unittest/llmapi/apps/utils.py index 783f6937bd53..25fb6da165e1 100644 --- a/tests/unittest/llmapi/apps/utils.py +++ b/tests/unittest/llmapi/apps/utils.py @@ -153,8 +153,7 @@ def make_server_with_custom_sampler_fixture(api_type: str) -> Callable: @pytest.fixture(scope='function') def server_with_custom_sampler(model_name: str, request: Any, backend: str, tmp_path: Path) -> RemoteOpenAIServer: - '''Fixture to launch a server (pytorch backend only) with a custom sampler configuration.''' - sampler_type = getattr(request, 'param', {}).get('sampler_type', "auto") + '''Fixture to launch a server (pytorch backend only) for sampling tests.''' if backend != 'pytorch': pytest.skip( f"Server with custom sampler is only supported for pytorch backend, skipping for {backend}" @@ -162,10 +161,7 @@ def server_with_custom_sampler(model_name: str, request: Any, backend: str, model_path = get_model_path(model_name) args = ['--backend', backend] temp_file_path = tmp_path / f'test_sampler_config_{request.node.name}.yaml' - extra_llm_api_options_dict = { - 'enable_chunked_prefill': True, - 'sampler_type': sampler_type - } + extra_llm_api_options_dict = {'enable_chunked_prefill': True} with temp_file_path.open('w') as f: yaml.dump(extra_llm_api_options_dict, f) args.extend(['--extra_llm_api_options', str(temp_file_path)]) diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index 2fadf2d0fd32..4a7560e43718 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -85,7 +85,7 @@ def test_collect_llm_api_config_allows_approved_string_converters_only(): # Path, which is dropped because it is not an allowlisted scalar, while # union_backend's allowlisted str is captured. No production telemetry field # is Union[str, Path]; the only real Union allowlist fields are - # Union[str, Enum] (sampler_type, load_format). See CR-E (declined). + # Union[str, Enum] (load_format). See CR-E (declined). class _StringConfig(StrictBaseModel): backend: Optional[str] = Field( default="pytorch", @@ -144,7 +144,7 @@ def test_sanitize_allowlist_is_value_fail_closed_for_non_scalars(): (bool/int/float/str) or None. Excluding Union-with-Any/Path from allowlist eligibility at the type level is therefore unnecessary for safety, and a coarse rule would also break legitimate Union[str, Enum] allowlist fields - such as sampler_type and load_format (verified captured elsewhere). + such as load_format (verified captured elsewhere). """ from tensorrt_llm.usage import llmapi_config @@ -626,20 +626,6 @@ class _C(StrictBaseModel): assert meta["unsafe_excluded"] is False -def test_collect_llm_api_config_captures_sampler_type_categorical(): - """sampler_type is a bounded Union[str, SamplerType] categorical allowlist.""" - args = TorchLlmArgs( - model="/customer/private/Llama", - skip_tokenizer_init=True, - sampler_type="TorchSampler", - ) - - config, meta = _loads_payloads(args) - - assert config["sampler_type"] == "TorchSampler" - assert meta["capture_succeeded"] is True - - def test_collect_llm_api_config_captures_transceiver_runtime_categorical(): """transceiver_runtime is a single Optional[Literal] categorical.